##############################################################################
#
#   Kennedy Institute of Rheumatology
#
#   $Id$
#
#   Copyright (C) 2015 Stephen Sansom
#
#   This program 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 2
#   of the License, or (at your option) any later version.
#
#   This program 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 this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
###############################################################################

"""===========================
Pipeline template
===========================

:Author: Stephen Sansom
:Release: $Id$
:Date: |today|
:Tags: Python


Overview
========

A pipeline designed to work with SMART-seq style single cell data
and bulk RNA-seq data. Single and paired end-data is supported.

This pipeline performs the follow tasks:

(1) [optional] Mapping of reads using hisat
      - paired (default) or single end fastq files are expected as the input
      - unstranded (default) or stranded fastq files are expected as the input
      - If data is already mapped position sorted, indexed BAM files can be
        provided instead.
      - See pipeline.ini and below for filename syntax guidance.

(2) Quantitation of gene expression
      - Ensembl geneset or supplied GTF (+/- spike in sequences)
      - Salmon is used to calculate TPM values
      - Cufflinks (cuffquant + cuffnorm) can be optionally run
      - featureCounts (from the subread package) is run for counting reads

(3) Calculation of post-mapping QC statistics
      - Picard CollectRnaSeqMetrics, AlignmentSummaryMetrics,
               EstimateLibraryComplexity, InsertSizeMetrics
      - Fraction spliced reads
      - Ratio spike-ins / genomic
      - Three prime bias
      - Numvers of genes detected

Usage
=====

See :ref:`PipelineSettingUp` and :ref:`PipelineRunning` on general
information how to use CGAT pipelines.

Configuration
-------------

The pipeline requires a configured :file:`pipeline.yml` file.
CGATReport report requires a :file:`conf.py` and optionally a
file:`cgatreport.ini` file (see :ref:`PipelineReporting`).

Default configuration files can be generated by executing:

   python <srcdir>/pipeline_scrnaseq.py config


Input files
-----------

* Fastqs

- The pipeline expects sequence data from each cell/sample in the form
  of single or paired-end fastq files to be present in the "input_dir"
  specificed in the pipeline.ini file (default = "data.dir")

- where there are multiple Fastq file per sample, the per-sample fastqs
  should be placed in directories which end with the pattern:
  (i) ".fastq.1.gz" - for paired end data
  (ii) ".fastq.gz" - for single end data

* BAMs

- location of directory containing the BAM files is specified in the
  "input_dir" variable in pipeline.ini (default = "data.dir")

* Naming

It is recommended that files are named according
to the following convention (for plate-based single-cell data):

source_condition_replicate_plate_row_column

An arbitrary number of fields can be specified, e.g. see the
pipeline.yml default:

    name_field_titles: source,condition,plate,row,column

Name fields *must be separated by underscores* (for SQL/R compatibility).

example fastq file name:

    mTEChi_wildtype_plate1_A_1.fastq.1.gz

If the pipeline is being used to compare BAM files from different mappers,
the mapper should be supplied as an extra name field, e.g.

  name_field_titles: source,condition,plate,row,column,mapper

    mTEChi_wildtype_plate1_A_1_gsnap.bam
    mTEChi_wildtype_plate1_A_1_gsnap.bam.bai

* Spike-in information

If spike-ins are used, the location of a table containing the per-cell
spike in copy numbers should be provided in a file specified in pipeline.yml

The expected structure (tab-delimited) is:

    gene_id|copies_per_cell
    ERCC-..|20324.7225

Note that the columns headers "gene_id" and "copies_per_cell" are required.


Requirements
------------

On top of the default CGAT setup, the pipeline requires the following
software to be in the path:

.. Add any additional external requirements such as 3rd party software
   or R modules below:

Requirements (TBC):

* Cufflinks
* Picard
* Hisat2
* subread
* R
* Salmon
* Kent tools

Pipeline output
===============

.. TBC

Glossary
========

.. glossary::


Code
====

"""
from ruffus import *

import sys
import shutil
import os
from pathlib import Path
import glob
import sqlite3

import pandas as pd
import numpy as np

from cgatcore import experiment as E
from cgatcore import pipeline as P
from cgatcore import database as DB

import PipelineScRnaseq as PipelineScRnaseq


# -------------------------- < parse parameters > --------------------------- #

# load options from the config file
PARAMS = P.get_parameters(
    ["%s/pipeline.yml" % os.path.splitext(__file__)[0],
     "../pipeline.yml",
     "pipeline.yml"])


# ----------------------- < pipeline configuration > ------------------------ #

if len(sys.argv) > 1:
    if(sys.argv[1] == "config") and __name__ == "__main__":
        sys.exit(P.main(sys.argv))

# Establish the location of module scripts for P.submit() functions
# set the location of the tenx code directory
if "code_dir" not in PARAMS.keys():
    code_dir = os.path.dirname(os.path.realpath(__file__))
else:
    raise ValueError("Could not set the location of the code directory")


# Set the database locations
DATABASE = PARAMS["database"]["file"]
ANN_DATABASE = PARAMS["annotations_database"]

# set the location of the scseq code directory
if "scseq_dir" not in PARAMS.keys():
    PARAMS["scseq_dir"] = Path(__file__).parents[1]
else:
    raise ValueError("Could not set the location of the scseq code directory")


# ########################################################################### #
# ############# Check sample files and prepare Sample table  ################ #
# ########################################################################### #

if PARAMS["input_type"].lower() == "fastq":
    suffix_pattern = "*.fastq.*gz"
elif PARAMS["input_type"].lower() == "bam":
    suffix_pattern = "*.bam"
else:
    raise ValueError("this pipeline only supports fastq or bam files")

SAMPLE_FILES = glob.glob(os.path.join(PARAMS["input_dir"], suffix_pattern))

# Check we have input
if len(SAMPLE_FILES) == 0:
    raise ValueError("No input files detected")

SAMPLE_IDS = set([os.path.basename(SF).split(".")[0] for SF in SAMPLE_FILES])
NAME_FIELD_TITLES = PARAMS["name_field_titles"].split(",")

# Check field names
if any([x in NAME_FIELD_TITLES for x in ("sample_id")]):
    raise ValueError('"sample_id" is reserved and cannot'
                     'be used as a name field title')

# Sanity check file names
for sample_id in SAMPLE_IDS:
    if len(sample_id.split("_")) != len(NAME_FIELD_TITLES):
        raise ValueError("%(sample_id)s does not have the expected"
                         " number of name fields (%(NAME_FIELD_TITLES)s)."
                         " Note that name fields must be separated with"
                         " underscores" % locals())

# Prepare sample table
SAMPLES = pd.DataFrame([dict(list(zip(["sample_id"] + NAME_FIELD_TITLES,
                                      [SAMPLE_ID] + SAMPLE_ID.split("_"))))
                        for SAMPLE_ID in SAMPLE_IDS])


# ########################################################################### #
# ######## Define endedness, strandedness and spike-in parameters ########### #
# ########################################################################### #

# Determine endedness
if str(PARAMS["paired"]).lower() in ("1", "true", "yes"):
    PAIRED = True
elif str(PARAMS["paired"]).lower() in ("0", "false", "no"):
    PAIRED = False
else:
    raise ValueError("Endedness not recognised")

# set options based on strandedness
STRAND = str(PARAMS["strandedness"]).lower()
if STRAND not in ("none", "forward", "reverse"):
    raise ValueError("Strand not recognised")

if STRAND == "none":
    CUFFLINKS_STRAND = "fr-unstranded"
    FEATURECOUNTS_STRAND = "0"
    PICARD_STRAND = "NONE"
    SALMON_STRAND = "U"

elif STRAND == "forward":
    if PAIRED:
        HISAT_STRAND = "FR"
    else:
        HISAT_STRAND = "F"
    SALMON_STRAND = "SF"
    CUFFLINKS_STRAND = "fr-secondstrand"
    FEATURECOUNTS_STRAND = "1"
    PICARD_STRAND = "FIRST_READ_TRANSCRIPTION_STRAND"

elif STRAND == "reverse":
    if PAIRED:
        HISAT_STRAND = "RF"
    else:
        HISAT_STRAND = "R"
    SALMON_STRAND = "SR"
    CUFFLINKS_STRAND = "fr-firststrand"
    FEATURECOUNTS_STRAND = "2"
    PICARD_STRAND = "SECOND_READ_TRANSCRIPTION_STRAND"

if PAIRED:
    SALMON_LIBTYPE = "I" + SALMON_STRAND
else:
    SALMON_LIBTYPE = SALMON_STRAND

SPIKES = PARAMS["spikein_present"]


# ########################################################################### #
# ################ Check annotation versions and genomes  ################### #
# ########################################################################### #

# Check that the genome source is set to "ucsc" or "ensembl"
GENOME_SOURCE = PARAMS["annotations_genome_source"].lower()

if not (GENOME_SOURCE == "ucsc" or GENOME_SOURCE == "ensembl"):
    raise ValueError('annotations_genome_source must be either'
                     ' "ucsc" or "ensembl" but "%(GENOME_SOURCE)s"'
                     ' was given' % locals())

# Check that given annotations and indices contain the expected
# Ensembl version number
ENSEMBL_VERSION = str(PARAMS["annotations_ensembl_version"])

if ENSEMBL_VERSION not in PARAMS["annotations_ensembl_geneset"]:
    raise ValueError("annotations geneset does not contain the"
                     " given ensembl version number")

if (PARAMS["annotations_geneset"] == "ensembl" and
        ENSEMBL_VERSION not in PARAMS["hisat_index"]):
    raise ValueError("hisat index does not contain the"
                     " given ensembl version number")

# Check that the given genome indices contain the expected
# genome name

GENOME_NAME = PARAMS["annotations_genome"]

if GENOME_NAME not in PARAMS["hisat_index"]:
    raise ValueError("hisat_index does not contain the"
                     " given genome name")

# Figure out which geneset to use for quantitation

if PARAMS["annotations_geneset"].lower() == "ensembl":
    QUANTITATION_GTF = PARAMS["annotations_ensembl_geneset"]
else:
    QUANTITATION_GTF = PARAMS["annotations_geneset"]

# Ensembl and UCSC have different contig naming conventions
# (UCSC prefixes contig names with "chr")
# Sanity check the contig patterns


@mkdir("preflight.checks.dir")
@files(os.path.join(PARAMS["annotations_genome_dir"],
                    PARAMS["annotations_genome"] + ".fasta"),
       "preflight.checks.dir/genome.contigs.txt")
def getGenomeContigs(infile, outfile):
    '''
    Extract contig names from genome FASTA file.
    '''

    statement = '''grep \\> %(infile)s
                   | sed 's/>//g'
                   | sort -u -T %(cluster_tmpdir)s
                   > %(outfile)s
                '''
    P.run(statement)


@mkdir("preflight.checks.dir")
@files(None,
       "preflight.checks.dir/hisat2.contigs.txt")
def getHisat2Contigs(infile, outfile):
    '''
    Extract contig names from HISAT2 index file.
    '''

    statement = '''hisat2-inspect -s %(hisat_index)s
                   | grep Sequence
                   | cut -f2
                   | sort -u -T %(cluster_tmpdir)s
                   > %(outfile)s
                '''

    P.run(statement)


@mkdir("preflight.checks.dir")
@files(QUANTITATION_GTF,
       "preflight.checks.dir/quantification.geneset.contigs.txt")
def getQuantitationGenesetContigs(infile, outfile):
    '''
    Extract contig names the GTF file used for quantitation.
    '''

    statement = '''zcat %(infile)s
                   | grep -v ^#
                   | cut -f1
                   | sort -u -T %(cluster_tmpdir)s
                   > %(outfile)s
                '''

    P.run(statement)


@mkdir("preflight.checks.dir")
@files(PARAMS["annotations_ensembl_geneset"],
       "preflight.checks.dir/ensembl.geneset.contigs.txt")
def getEnsemblGenesetContigs(infile, outfile):
    '''
    Extract contig names from the ENSEMBL GTF file of gene annotation.
    '''

    statement = '''zcat %(infile)s
                   | grep -v ^#
                   | cut -f1
                   | sort -u -T %(cluster_tmpdir)s
                   > %(outfile)s
                '''

    P.run(statement)


@merge([getGenomeContigs, getHisat2Contigs,
        getQuantitationGenesetContigs,
        getEnsemblGenesetContigs],
       "preflight.checks.dir/contig.report.txt")
def checkContigs(infiles, outfile):
    '''
    Check that all the annotations and indices
    share a common set of contigs.
    '''

    contig_sets = {}
    contig_ns = []

    report = open(outfile, "w")

    for infile in infiles:

        set_name = os.path.basename(infile).split(".")[0]
        contig_sets[set_name] = []

        with open(infile, "r") as contig_file:

            for line in contig_file:

                contig = line.strip()

                if contig.startswith("chr"):

                    if GENOME_SOURCE == "ucsc":
                        pass
                    else:
                        raise ValueError("ensembl contigs are not expected"
                                         " to contain chr pattern"
                                         " in %(infile)s" % locals())

                else:
                    if GENOME_SOURCE == "ensembl":
                        pass
                    else:
                        raise ValueError("ucsc chr contig pattern not found"
                                         " in %(infile)s" % locals())

                contig_sets[set_name].append(contig)

        n_contigs = len(contig_sets[set_name])
        contig_ns.append(n_contigs)

        if n_contigs == 0:
            raise ValueError("No contigs found in %(infile)s" % locals())

        report.write("#%(set_name)s: found %(n_contigs)i contigs\n" % locals())

    begin = True

    for contig_set, contigs in contig_sets.items():

        if begin:
            common = set(contigs)
            begin = False

        else:
            common = common.intersection(set(contigs))

    n_common = len(common)
    if n_common != np.max([np.min(contig_ns), 1]):

        raise ValueError("No common contigs found")

    report.write("Found %(n_common)i common contigs:\n" % locals())

    for contig in common:
        report.write(contig + "\n")

    report.close()


# ---------------------- < specific pipeline tasks > ------------------------ #

# ########################################################################### #
# #################### Generate salmon index ################################ #
# ########################################################################### #

SALMON_INDEX_NAME = ".".join(
    [os.path.basename(QUANTITATION_GTF[:-len(".gtf.gz")]),
     PARAMS["salmon_index_type"],
     str(PARAMS["salmon_index_k"])])
SALMON_INDEX = os.path.join("annotations.dir", SALMON_INDEX_NAME)


@active_if(PARAMS["input_type"] == "fastq")
@follows(mkdir("annotations.dir"), checkContigs)
@files(QUANTITATION_GTF,
       "annotations.dir/salmon_build.log")
def generateSalmonIndex(infile, outfile):
    '''Generate salmon index from annotation gtf. '''
    genome_fasta = os.path.join(PARAMS["annotations_genome_dir"],
                                PARAMS["annotations_genome_fasta"])
    index_folder = SALMON_INDEX
    job_memory = PARAMS["salmon_index_memory"]
    statement = '''fasta_out=`mktemp -p %(cluster_tmpdir)s`;
            gtf=`mktemp -p %(cluster_tmpdir)s`;
            zcat %(infile)s > $gtf ;
             gffread
            -w $fasta_out
            -g %(genome_fasta)s
            $gtf ;
            salmon index
            -t $fasta_out
            -i %(index_folder)s
            %(salmon_index_options)s
            --type %(salmon_index_type)s
            -k %(salmon_index_k)s
            &> %(outfile)s '''
    P.run(statement)


# ########################################################################### #
# #################### (1) Read Mapping (optional) ########################## #
# ########################################################################### #


if PAIRED:
        fastq_pattern = "*.fastq.1.gz"
else:
        fastq_pattern = "*.fastq.gz"

if STRAND != "none":
    HISAT_STRAND_PARAM = "--rna-strandness %s" % HISAT_STRAND
else:
    HISAT_STRAND_PARAM = ""

HISAT_THREADS = PARAMS["hisat_threads"]
HISAT_MEMORY = str(int(PARAMS["hisat_total_mb_memory"]) //
                   int(HISAT_THREADS)) + "M"


@follows(mkdir("hisat.dir/first.pass.dir"), checkContigs)
@transform(glob.glob(os.path.join(PARAMS["input_dir"], fastq_pattern)),
           regex(r".*/(.*).fastq.*.gz"),
           r"hisat.dir/first.pass.dir/\1.novel.splice.sites.txt.gz")
def hisatFirstPass(infile, outfile):
    '''
    Run a first hisat pass to identify novel splice sites.
    '''

    reads_one = infile

    if os.path.isdir(reads_one):
        reads_one = glob.glob(os.path.join(reads_one, fastq_pattern))

    else:
        reads_one = [reads_one]

    index = PARAMS["hisat_index"]
    log = outfile + ".log"
    out_name = outfile[:-len(".gz")]

    # queue options
    to_cluster = True  # this is the default
    job_threads = HISAT_THREADS
    job_memory = HISAT_MEMORY

    if PAIRED:
        reads_two = [x[:-len(".1.gz")] + ".2.gz" for x in reads_one]
        fastq_input = "-1 " + ",".join(reads_one) +\
                      " -2 " + ",".join(reads_two)
    else:
        fastq_input = "-U " + ",".join(reads_one)

    hisat_strand_param = HISAT_STRAND_PARAM

    statement = '''%(hisat_executable)s
                        -x %(index)s
                        %(fastq_input)s
                        --threads %(job_threads)s
                        --novel-splicesite-outfile %(out_name)s
                        %(hisat_strand_param)s
                        %(hisat_options)s
                        -S /dev/null
                        &> %(log)s;
                    gzip %(out_name)s
                '''

    P.run(statement)


@follows(mkdir("annotations.dir"))
@merge(hisatFirstPass, "annotations.dir/novel.splice.sites.hisat.txt")
def novelHisatSpliceSites(infiles, outfile):
    '''
    Collect the novel splice sites into a single file.
    '''

    junction_files = " ".join(infiles)

    statement = '''zcat %(junction_files)s
                   | sort -k1,1 -T %(cluster_tmpdir)s
                   | uniq
                   > %(outfile)s
                '''

    P.run(statement)


@transform(glob.glob(os.path.join(PARAMS["input_dir"], fastq_pattern)),
           regex(r".*/(.*).fastq.*.gz"),
           add_inputs(novelHisatSpliceSites),
           r"hisat.dir/\1.bam")
def hisatAlignments(infiles, outfile):
    '''
    Align reads using HISAT with known and novel junctions.
    '''

    reads_one, novel_splice_sites = infiles

    if os.path.isdir(reads_one):
        reads_one = glob.glob(os.path.join(reads_one, fastq_pattern))
    else:
        reads_one = [reads_one]

    index = PARAMS["hisat_index"]
    log = outfile + ".log"

    to_cluster = True
    job_threads = HISAT_THREADS
    job_memory = HISAT_MEMORY

    if PAIRED:
        reads_two = [x[:-len(".1.gz")] + ".2.gz" for x in reads_one]
        fastq_input = "-1 " + ",".join(reads_one) +\
                      " -2 " + ",".join(reads_two)
    else:
        fastq_input = "-U " + ",".join(reads_one)

    hisat_strand_param = HISAT_STRAND_PARAM

    statement = '''sort_sam=`mktemp -p %(cluster_tmpdir)s`;
                   %(hisat_executable)s
                      -x %(index)s
                      %(fastq_input)s
                      --threads %(job_threads)s
                      --novel-splicesite-infile %(novel_splice_sites)s
                      %(hisat_strand_param)s
                      %(hisat_options)s
                   2> %(log)s
                   | samtools view - -bS
                   | samtools sort - -T $sort_sam -o %(outfile)s >>%(log)s;
                   samtools index %(outfile)s;
                   rm $sort_sam;
                 '''

    P.run(statement)


@follows(hisatAlignments)
def mapping():
    '''
    Mapping target.
    '''
    pass


# ########################################################################### #
# ########### Collect BAMs from mapping functions or inputs  ################ #
# ########################################################################### #

if PARAMS["input_type"] == "fastq":
    collectBAMs = hisatAlignments
    fastqMode = True

elif PARAMS["input_type"] == "bam":
    collectBAMs = glob.glob(os.path.join(PARAMS["input_dir"], "*.bam"))
    fastqMode = False

else:
    raise ValueError('Input type must be either "fastq" or "bam"')


# ########################################################################### #
# ################ (2) Quantitation of gene expression #################### #
# ########################################################################### #

# ------------------------- Geneset Definition ------------------------------ #

@follows(mkdir("annotations.dir"), checkContigs)
@files(QUANTITATION_GTF,
       "annotations.dir/quantitation.geneset.gtf.gz")
def prepareQuantitationGenesetGTF(infile, outfile):
    '''
    Preparation of the GTF file used for quantitation:
    Spike-in GTF entries are optionally appended
    to the reference annotation.
    '''

    outname = outfile[:-len(".gz")]

    geneset_stat = '''zcat %(infile)s
                     > %(outname)s;
                  '''
    if SPIKES:
        spikein_geneset = PARAMS["spikein_geneset"]
        spikein_stat = '''cat %(spikein_geneset)s >> %(outname)s;
                        '''
    else:
        spikein_stat = ''

    statement = geneset_stat + spikein_stat + "gzip %(outname)s"

    P.run(statement)


@follows(mkdir("annotations.dir"))
@files(PARAMS["annotations_ensembl_geneset"],
       "annotations.dir/ensembl.geneset.flat.gz")
def prepareEnsemblGenesetFlat(infile, outfile):
    '''
    Prepare a flat version of the geneset
    for the Picard CollectRnaSeqMetrics module.
    '''

    statement = '''gtfToGenePred
                    -genePredExt
                    -geneNameAsName2
                    -ignoreGroupsWithoutExons
                    %(infile)s
                    /dev/stdout |
                    awk 'BEGIN { OFS="\\t"}
                         {print $12, $1, $2, $3, $4, $5, $6, $7, $8, $9, $10}'
                    | gzip -c
                    > %(outfile)s
                 '''

    P.run(statement)


@follows(mkdir("annotations.dir"))
@follows(prepareQuantitationGenesetGTF)
@files(PARAMS["annotations_ensembl_geneset"],
       "annotations.dir/transcript_info.txt.gz")
def tabulateTranscriptInfoFromGTF(infile, outfile):
    '''
    Tabulate selected transcript-level metadata from a GTF.
    '''
    job_memory = "10G"

    extract_fields = ",".join(['gene_id', 'transcript_id',
                               'gene_biotype', 'transcript_biotype',
                               'gene_name'])

    log_file = outfile.replace("txt.gz", "log")

    statement = '''Rscript %(scseq_dir)s/R/tabulate_transcript_information.R
                   --gtf=%(infile)s
                   --fields=%(extract_fields)s
                   --outfile=%(outfile)s
                   &> %(log_file)s
                '''
    P.run(statement)


@transform(tabulateTranscriptInfoFromGTF,
           suffix(".txt.gz"),
           ".load")
def loadEnsemblAnnotations(infile, outfile):
    '''
    Load the annotations for salmon into the project database.
    '''

    # will use ~15G RAM
    P.load(infile, outfile, options='-i "gene_id" -i "transcript_id"')


@transform(tabulateTranscriptInfoFromGTF,
           regex("(.*)/.*"),
           r"\1/tx2gene.txt")
def tx2gene(infile, outfile):
    '''
    Preparation of transcript to gene map for use by Salmon and tximport.
    '''

    tx_info = pd.read_csv(infile, sep="\t", header=0)
    tx_info = tx_info[["transcript_id", "gene_id"]]
    tx_info = tx_info[pd.notnull(tx_info["transcript_id"])]
    tx_info = tx_info.drop_duplicates()
    tx_info.to_csv(outfile, sep="\t", header=None, index=False)

    # the above is expected to be equivalent too:
    #
    # geneset_stat = '''zcat %(geneset)s
    #  | grep transcript_id
    #  | sed 's/.*gene_id "\([^"]*\)".*transcript_id "\([^"]*\)".*/\2\t\1/g'
    #  | sort -u
    #  > %(outfile)s;
    #  '''

    if SPIKES:
        spikein_tx2db = PARAMS["spikein_tx2gene"]
        spikein_stat = '''cat %(spikein_tx2gene)s
                          >> %(outfile)s;
                       '''
        statement = spikein_stat

        P.run(statement)


# Figure out which source to use for mappings between transcript_id and gene_id
if PARAMS["salmon_tx2gene"].lower() == "ensembl":
    TX2GENE = tx2gene
else:
    TX2GENE = PARAMS["salmon_tx2gene"]


@follows(mkdir("annotations.dir"), checkContigs)
@files(QUANTITATION_GTF,
       "annotations.dir/quantitation.geneset.gtf.gz")
def prepareTranscript(infile, outfile):
    '''
    Preparation of the GTF file used for quantitation:
    Spike-in GTF entries are optionally appended
    to the reference annotation.
    '''

    outname = outfile[:-len(".gz")]

    geneset_stat = '''zcat %(infile)s
                     > %(outname)s;
                  '''
    if SPIKES:
        spikein_geneset = PARAMS["spikein_geneset"]
        spikein_stat = '''cat %(spikein_geneset)s >> %(outname)s;
                        '''
    else:
        spikein_stat = ''

    statement = geneset_stat + spikein_stat + "gzip %(outname)s"

    P.run(statement)


@follows(prepareQuantitationGenesetGTF,
         prepareEnsemblGenesetFlat,
         loadEnsemblAnnotations,
         tx2gene)
def annotations():
    '''
    Annotation target.
    '''
    pass


# ----------------------------- Read Counting ------------------------------- #

@follows(mkdir("featureCounts.dir"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           add_inputs(prepareQuantitationGenesetGTF),
           r"featureCounts.dir/\1.counts.gz")
def featureCounts(infiles, outfile):
    '''
    Run featureCounts.
    Note that we first need to change directory to a scratch location because
    the current dir is hard coded as the temp dir!!
    '''

    bamfile, geneset = [os.path.abspath(x) for x in infiles]
    outfile_name = os.path.abspath(outfile)

    # set featureCounts options
    featurecounts_strand = FEATURECOUNTS_STRAND

    if PAIRED:
        paired_options = "-p"
    else:
        paired_options = ""

    if PARAMS["featurecounts_options"] is None:
        featurecounts_options = ""
    else:
        featurecounts_options = PARAMS["featurecounts_options"]

    job_threads = PARAMS["featurecounts_threads"]

    statement = '''cd %(cluster_tmpdir)s;
                   gtf=`mktemp -p %(cluster_tmpdir)s`;
                   counts=`mktemp -p %(cluster_tmpdir)s`;
                   zcat %(geneset)s > $gtf;
                   featureCounts
                        -a $gtf
                        -o $counts
                        -s %(featurecounts_strand)s
                        -T %(featurecounts_threads)s
                        %(featurecounts_options)s
                        %(paired_options)s
                        %(bamfile)s;
                        cut -f1,7 $counts
                        | grep -v "#" | grep -v "Geneid"
                        | gzip -c > %(outfile_name)s;
                        rm $gtf;
                        rm $counts;
                 '''

    P.run(statement)


@merge(featureCounts,
       "featureCounts.dir/featurecounts.load")
def loadFeatureCounts(infiles, outfile):
    '''
    Combine and load count data in the project database.
    '''

    P.concatenate_and_load(infiles, outfile,
                           regex_filename=".*/(.*).counts.gz",
                           has_titles=False,
                           cat="track",
                           header="track,gene_id,counts",
                           options='-i "gene_id"',
                           job_memory=PARAMS["sql_himem"])


@files(loadFeatureCounts,
       "featureCounts.dir/featurecounts_counts.txt")
def featurecountsGeneCounts(infile, outfile):
    '''
    Prepare a gene-by-sample table of featureCounts counts.
    '''

    table = P.to_table(infile)
    con = sqlite3.connect(PARAMS["database_file"])
    c = con.cursor()

    sql = '''select track, gene_id, counts
             from %(table)s t
          ''' % locals()

    df = pd.read_sql(sql, con)
    df = df.pivot("gene_id", "track", "counts")
    df.to_csv(outfile, sep="\t", index=True, index_label="gene_id")


@transform(featurecountsGeneCounts,
           suffix(".txt"),
           ".load")
def loadFeaturecountsTables(infile, outfile):
    '''
    Load the gene-by-sample matrix of count data in the project database.
    '''

    P.load(infile, outfile, options='-i "gene_id"')


# ---------------------- Salmon TPM calculation ----------------------------- #

@active_if(fastqMode)
@follows(mkdir("salmon.dir"), generateSalmonIndex)
@transform(glob.glob(os.path.join(PARAMS["input_dir"], fastq_pattern)),
           regex(r".*/(.*).fastq.*.gz"),
           add_inputs(TX2GENE),
           r"salmon.dir/\1.log")
def salmon(infiles, outfile):
    '''
    Per sample quantitation using salmon.
    '''

    reads_one, tx2gene = infiles
    salmon_index = SALMON_INDEX
    outname = outfile[:-len(".log")]

    if os.path.isdir(reads_one):
        reads_one = glob.glob(os.path.join(reads_one, fastq_pattern))

    else:
        reads_one = [reads_one]

    if PAIRED:
        reads_two = [x[:-len(".1.gz")] + ".2.gz" for x in reads_one]
        fastq_input = "-1 " + " ".join(reads_one) +\
                      " -2 " + " ".join(reads_two)

    else:
        fastq_input = "-r " + " ".join(reads_one)

    if not PARAMS['salmon_params'] is None:
        salmon_options = PARAMS['salmon_params']
    else:
        salmon_options = ''

    salmon_libtype = SALMON_LIBTYPE
    salmon_params = PARAMS["salmon_params"]
    job_threads = PARAMS["salmon_threads"]

    statement = '''salmon quant -i %(salmon_index)s
                                -p %(job_threads)s
                                -g %(tx2gene)s
                                %(salmon_options)s
                                -l %(salmon_libtype)s
                                %(fastq_input)s
                                -o %(outname)s
                    &> %(outfile)s;
              '''

    P.run(statement)


@active_if(fastqMode)
@merge(salmon, "salmon.dir/salmon.transcripts.load")
def loadSalmonTranscriptQuant(infiles, outfile):
    '''
    Load the salmon transcript-level results.
    '''

    tables = [x.replace(".log", "/quant.sf") for x in infiles]

    P.concatenate_and_load(tables, outfile,
                           regex_filename=".*/(.*)/quant.sf",
                           cat="sample_id",
                           options="-i Name -i sample_id",
                           job_memory=PARAMS["sql_himem"])


@active_if(fastqMode)
@merge(salmon, "salmon.dir/salmon.genes.load")
def loadSalmonGeneQuant(infiles, outfile):
    '''
    Load the salmon gene-level results.
    '''

    tables = [x.replace(".log", "/quant.genes.sf") for x in infiles]

    P.concatenate_and_load(tables, outfile,
                           regex_filename=".*/(.*)/quant.genes.sf",
                           cat="sample_id",
                           options="-i Name -i sample_id",
                           job_memory=PARAMS["sql_himem"])


@active_if(fastqMode)
@jobs_limit(1)
@transform([loadSalmonTranscriptQuant,
            loadSalmonGeneQuant],
           regex(r"(.*)/(.*).load"),
           r"\1/\2.tpms.txt")
def salmonTPMs(infile, outfile):
    '''
    Prepare a wide table of salmon TPMs (samples x transcripts|genes).
    '''

    table = P.to_table(infile)

    if "transcript" in table:
        id_name = "transcript_id"
    elif "gene" in table:
        id_name = "gene_id"
    else:
        raise ValueError("Unexpected Salmon table name")

    con = sqlite3.connect(PARAMS["database_file"])
    c = con.cursor()

    sql = '''select sample_id, Name %(id_name)s, TPM tpm
             from %(table)s
          ''' % locals()

    df = pd.read_sql(sql, con)

    df = df.pivot(id_name, "sample_id", "tpm")
    df.to_csv(outfile, sep="\t", index=True, index_label=id_name)


@active_if(fastqMode)
@transform(salmonTPMs,
           suffix(".txt"),
           ".load")
def loadSalmonTPMs(infile, outfile):
    '''
    Load a wide table of salmon TPMs in the project database.
    '''

    if "transcript" in infile:
        id_name = "transcript_id"
    elif "gene" in infile:
        id_name = "gene_id"
    else:
        raise ValueError("Unexpected Salmon table name")

    opts = "-i " + id_name

    P.load(infile, outfile, options=opts,
           job_memory=PARAMS["sql_himem"])


# For summarisation of Salmon counts, downstream use of
# tximport is recommended.
# For this the annotations.dir/tx2gene.txt file provides
# the transcript -> gene mappings.

# -------------------- FPKM (Cufflinks) quantitation------------------------- #

# Optional cufflinks based-quantitation, e.g.
#
# "make loadCuffNormClassic"
# "make loadCuffNormUQ"

@follows(mkdir("cuffquant.dir"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           add_inputs(prepareQuantitationGenesetGTF),
           r"cuffquant.dir/\1.log")
def cuffQuant(infiles, outfile):
    '''
    Per sample quantitation using cuffquant.
    '''

    bam_file, geneset = infiles

    # because the output of cuffquant is always "abundances.cxb"
    # a unique output directory for each sample is required
    output_dir = outfile[:-len(".log")]

    job_threads = PARAMS["cufflinks_cuffquant_threads"]  # 4

    to_cluster = True

    genome_multifasta = os.path.join(PARAMS["annotations_genome_dir"],
                                     PARAMS["annotations_genome"]+".fasta")

    cufflinks_strand = CUFFLINKS_STRAND

    statement = '''gtf=`mktemp -p %(cluster_tmpdir)s`;
                   zcat %(geneset)s > $gtf;
                   cuffquant
                           --output-dir %(output_dir)s
                           --num-threads %(job_threads)s
                           --multi-read-correct
                           --library-type %(cufflinks_strand)s
                           --no-effective-length-correction
                           --max-bundle-frags 2000000
                           --max-mle-iterations 10000
                           --verbose
                           --frag-bias-correct %(genome_multifasta)s
                            $gtf %(bam_file)s >& %(outfile)s;
                    rm $gtf;
                '''

    P.run(statement)


@follows(mkdir("cuffnorm.classic.dir"), cuffQuant)
@merge([prepareQuantitationGenesetGTF, cuffQuant],
       "cuffnorm.classic.dir/cuffnorm_classic.log")
def cuffNormClassic(infiles, outfile):
    '''
    Calculate classic FPKMs using cuffNorm
    for copy number estimation.
    '''

    cxb_files = " ".join([f[:-len(".log")] + "/abundances.cxb"
                          for f in infiles[1:]])

    label_str = ",".join([os.path.basename(f)[:-len(".log")]
                          for f in infiles[1:]])

    # parse the infiles
    geneset = infiles[0]

    # get the output directory and sample labels
    output_dir = os.path.dirname(outfile)

    PipelineScRnaseq.runCuffNorm(geneset, cxb_files, label_str,
                                 output_dir, outfile,
                                 library_type=CUFFLINKS_STRAND,
                                 normalisation="classic-fpkm", hits="total")


@transform(cuffNormClassic,
           suffix(".log"),
           ".load")
def loadCuffNormClassic(infile, outfile):
    '''
    Load the fpkm table from cuffnorm into the project database.
    '''

    fpkm_table = os.path.dirname(infile) + "/genes.fpkm_table"

    to_cluster = True

    P.load(fpkm_table, outfile,
           options='-i "tracking_id"')


@follows(mkdir("cuffnorm.uq.dir"), cuffQuant)
@merge([prepareQuantitationGenesetGTF, cuffQuant],
       "cuffnorm.uq.dir/cuffnorm_uq.log")
def cuffNormUQ(infiles, outfile):
    '''
    Calculate upper quartile (UQ) normalised FPKMs using cuffNorm.
    '''

    # parse the infiles
    geneset = infiles[0]

    cxb_path = os.path.dirname(infiles[1])
    cxb_name = "abundances.cxb"

    # Group replicate samples
    replicate_field = PARAMS["cufflinks_replicate_field"]

    if replicate_field:
        if replicate_field not in NAME_FIELD_TITLES:
            raise ValueError("cufflinks replicate field not in field titles")

        key = [T for T in NAME_FIELD_TITLES if T.lower() != replicate_field]
        agg = SAMPLES.groupby(key)

        labels = []
        cxb_groups = []
        for group, indices in agg.groups.items():

            labels.append("_".join(group))

            group_cxb_files = [os.path.join(cxb_path, S, cxb_name)
                               for S in
                               SAMPLES.ix[indices]["sample_id"].values]

            cxb_groups.append(",".join(group_cxb_files))

        cxb_files = " ".join(cxb_groups)

    else:

        sample_ids = SAMPLES["sample_id"].values

        cxb_files = " ".join([os.path.join(cxb_path, S, cxb_name)
                              for S in sample_ids])

        labels = sample_ids

    # get the output directory and sample labels
    output_dir = os.path.dirname(outfile)
    label_str = ",".join(labels)

    standards = PARAMS["cufflinks_standards"]

    PipelineScRnaseq.runCuffNorm(geneset, cxb_files, label_str,
                                 output_dir, outfile,
                                 library_type=CUFFLINKS_STRAND,
                                 standards_file=standards,
                                 normalisation="quartile", hits="compatible")


@transform(cuffNormUQ,
           suffix(".log"),
           ".load")
def loadCuffNormUQ(infile, outfile):
    '''
    Load the fpkm table from cuffnorm into the project database.
    '''

    fpkm_table = os.path.dirname(infile) + "/genes.fpkm_table"

    P.load(fpkm_table, outfile,
           options='-i "tracking_id"')


# ---------------------- Copynumber estimation ------------------------------ #

# Copy number estimation based on spike-in sequences and Salmon TPMs.
if PARAMS["spikein_estimate_copy_numbers"] is True:
    run_copy_number_estimation = True
else:
    run_copy_number_estimation = False


@active_if(run_copy_number_estimation)
@follows(mkdir("copy.number.dir"), loadSalmonTPMs)
@files("salmon.dir/salmon.genes.tpms.txt",
       "copy.number.dir/copy_numbers.txt")
def estimateCopyNumber(infile, outfile):
    '''
    Estimate copy numbers based on standard
    curves constructed from the spike-ins.
    '''

    statement = '''Rscript %(scseq_dir)s/R/calculate_copy_number.R
                   --spikeintable=%(spikein_copy_numbers)s
                   --spikeidcolumn=gene_id
                   --spikecopynumbercolumn=copies_per_cell
                   --exprstable=%(infile)s
                   --exprsidcolumn=gene_id
                   --outfile=%(outfile)s
                '''
    P.run(statement)


@active_if(run_copy_number_estimation)
@transform(estimateCopyNumber,
           suffix(".txt"),
           ".load")
def loadCopyNumber(infile, outfile):
    '''
    Load the copy number estimations to the project database.
    '''

    P.load(infile, outfile, options='-i "gene_id"')


# ----------------------- Quantitation target ------------------------------ #

@follows(annotations, loadFeatureCounts, loadFeaturecountsTables,
         loadSalmonTPMs, loadCopyNumber)
def quantitation():
    '''
    Quantitation target.
    '''
    pass


# ########################################################################### #
# ################ (3) Post-mapping Quality Control ######################### #
# ########################################################################### #

# -------------------- Set generic Picard options --------------------------- #

PICARD_THREADS = PARAMS["picard_threads"]
PICARD_MEMORY = str(int(PARAMS["picard_total_mb_memory"]) //
                    int(PICARD_THREADS)) + "M"


# ------------------- Picard: CollectRnaSeqMetrics -------------------------- #

@follows(mkdir("qc.dir/rnaseq.metrics.dir/"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           add_inputs(prepareEnsemblGenesetFlat),
           r"qc.dir/rnaseq.metrics.dir/\1.rnaseq.metrics")
def collectRnaSeqMetrics(infiles, outfile):
    '''
    Run Picard CollectRnaSeqMetrics on the bam files.
    '''

    bam_file, geneset_flat = infiles

    if PARAMS["picard_collectrnaseqmetrics_options"]:
        picard_options = PARAMS["picard_collectrnaseqmetrics_options"]
    else:
        picard_options = ""

    validation_stringency = PARAMS["picard_validation_stringency"]

    job_threads = PICARD_THREADS
    job_memory = PICARD_MEMORY

    coverage_out = outfile[:-len(".metrics")] + ".cov.hist"
    chart_out = outfile[:-len(".metrics")] + ".cov.pdf"

    picard_strand = PICARD_STRAND

    statement = '''picard_out=`mktemp -p %(cluster_tmpdir)s`;
                   CollectRnaSeqMetrics
                   I=%(bam_file)s
                   REF_FLAT=%(geneset_flat)s
                   O=$picard_out
                   CHART=%(chart_out)s
                   STRAND_SPECIFICITY=%(picard_strand)s
                   VALIDATION_STRINGENCY=%(validation_stringency)s
                   %(picard_options)s;
                   grep . $picard_out | grep -v "#" | head -n2
                   > %(outfile)s;
                   grep . $picard_out
                   | grep -A 102 "## HISTOGRAM"
                   | grep -v "##"
                   > %(coverage_out)s;
                   rm $picard_out;
                '''

    P.run(statement)


@merge(collectRnaSeqMetrics,
       "qc.dir/qc_rnaseq_metrics.load")
def loadCollectRnaSeqMetrics(infiles, outfile):
    '''
    Load the metrics to the db.
    '''

    P.concatenate_and_load(infiles, outfile,
                           regex_filename=".*/.*/(.*).rnaseq.metrics",
                           cat="sample_id",
                           options='-i "sample_id"')


# --------------------- Three prime bias analysis --------------------------- #

@transform(collectRnaSeqMetrics,
           suffix(".rnaseq.metrics"),
           ".three.prime.bias")
def threePrimeBias(infile, outfile):
    '''
    Compute a sensible three prime bias metric
    from the picard coverage histogram.
    '''

    coverage_histogram = infile[:-len(".metrics")] + ".cov.hist"

    df = pd.read_csv(coverage_histogram, sep="\t")

    x = "normalized_position"
    cov = "All_Reads.normalized_coverage"

    three_prime_coverage = np.mean(df[cov][(df[x] > 70) & (df[x] < 90)])
    transcript_body_coverage = np.mean(df[cov][(df[x] > 20) & (df[x] < 90)])
    bias = three_prime_coverage / transcript_body_coverage

    with open(outfile, "w") as out_file:
        out_file.write("three_prime_bias\n")
        out_file.write("%.2f\n" % bias)


@merge(threePrimeBias,
       "qc.dir/qc_three_prime_bias.load")
def loadThreePrimeBias(infiles, outfile):
    '''
    Load the metrics in the project database.
    '''

    P.concatenate_and_load(infiles, outfile,
                           regex_filename=".*/.*/(.*).three.prime.bias",
                           cat="sample_id",
                           options='-i "sample_id"')


# ----------------- Picard: EstimateLibraryComplexity ----------------------- #


@active_if(PAIRED)
@follows(mkdir("qc.dir/library.complexity.dir/"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           r"qc.dir/library.complexity.dir/\1.library.complexity")
def estimateLibraryComplexity(infile, outfile):
    '''
    Run Picard EstimateLibraryComplexity on the BAM files.
    '''

    if PARAMS["picard_estimatelibrarycomplexity_options"]:
        picard_options = PARAMS["picard_estimatelibrarycomplexity_options"]
    else:
        picard_options = ""

    validation_stringency = PARAMS["picard_validation_stringency"]

    job_threads = PICARD_THREADS
    job_memory = PICARD_MEMORY

    statement = '''picard_out=`mktemp -p %(cluster_tmpdir)s`;
                   EstimateLibraryComplexity
                   I=%(infile)s
                   O=$picard_out
                   VALIDATION_STRINGENCY=%(validation_stringency)s
                   %(picard_options)s;
                   grep . $picard_out | grep -v "#" | head -n2
                   > %(outfile)s;
                   rm $picard_out;
                '''

    P.run(statement)


@active_if(PAIRED)
@merge(estimateLibraryComplexity,
       "qc.dir/qc_library_complexity.load")
def loadEstimateLibraryComplexity(infiles, outfile):
    '''
    Load the complexity metrics to a single table in the project database.
    '''

    if PAIRED:
        P.concatenate_and_load(infiles, outfile,
                               regex_filename=".*/.*/(.*).library.complexity",
                               cat="sample_id",
                               options='-i "sample_id"')
    else:
        statement = '''echo "Not compatible with SE data"
                       > %(outfile)s'''
        P.run(statement)


# ------------------- Picard: AlignmentSummaryMetrics ----------------------- #

@follows(mkdir("qc.dir/alignment.summary.metrics.dir/"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           (r"qc.dir/alignment.summary.metrics.dir"
            r"/\1.alignment.summary.metrics"))
def alignmentSummaryMetrics(infile, outfile):
    '''
    Run Picard AlignmentSummaryMetrics on the bam files.
    '''

    picard_options = PARAMS["picard_alignmentsummarymetric_options"]
    validation_stringency = PARAMS["picard_validation_stringency"]

    job_threads = PICARD_THREADS
    job_memory = PICARD_MEMORY

    reference_sequence = os.path.join(PARAMS["annotations_genome_dir"],
                                      PARAMS["annotations_genome"] + ".fasta")

    statement = '''picard_out=`mktemp -p %(cluster_tmpdir)s`;
                   CollectAlignmentSummaryMetrics
                   I=%(infile)s
                   O=$picard_out
                   REFERENCE_SEQUENCE=%(reference_sequence)s
                   VALIDATION_STRINGENCY=%(validation_stringency)s
                   %(picard_options)s;
                   grep . $picard_out | grep -v "#"
                   > %(outfile)s;
                   rm $picard_out;
                '''

    P.run(statement)


@merge(alignmentSummaryMetrics,
       "qc.dir/qc_alignment_summary_metrics.load")
def loadAlignmentSummaryMetrics(infiles, outfile):
    '''
    Load the complexity metrics to a single table in the project database.
    '''

    P.concatenate_and_load(
        infiles, outfile,
        regex_filename=".*/.*/(.*).alignment.summary.metrics",
        cat="sample_id",
        options='-i "sample_id"')


# ------------------- Picard: InsertSizeMetrics ----------------------- #

@follows(mkdir("qc.dir/insert.size.metrics.dir/"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           [(r"qc.dir/insert.size.metrics.dir"
             r"/\1.insert.size.metrics.summary"),
            (r"qc.dir/insert.size.metrics.dir"
             r"/\1.insert.size.metrics.histogram")])
def insertSizeMetricsAndHistograms(infile, outfiles):
    '''
    Run Picard InsertSizeMetrics on the BAM files to
    collect summary metrics and histograms.'''

    if PAIRED:
        picard_summary, picard_histogram = outfiles
        picard_histogram_pdf = picard_histogram + ".pdf"

        if PARAMS["picard_insertsizemetric_options"]:
            picard_options = PARAMS["picard_insertsizemetric_options"]
        else:
            picard_options = ""

        job_threads = PICARD_THREADS
        job_memory = PICARD_MEMORY

        validation_stringency = PARAMS["picard_validation_stringency"]
        reference_sequence = os.path.join(PARAMS["annotations_genome_dir"],
                                          PARAMS["annotations_genome"] +
                                          ".fasta")

        statement = '''picard_out=`mktemp -p %(cluster_tmpdir)s`;
                       CollectInsertSizeMetrics
                       I=%(infile)s
                       O=$picard_out
                       HISTOGRAM_FILE=%(picard_histogram_pdf)s
                       VALIDATION_STRINGENCY=%(validation_stringency)s
                       REFERENCE_SEQUENCE=%(reference_sequence)s
                       %(picard_options)s;
                       grep "MEDIAN_INSERT_SIZE" -A 1 $picard_out
                       > %(picard_summary)s;
                       sed -e '1,/## HISTOGRAM/d' $picard_out
                       > %(picard_histogram)s;
                       rm $picard_out;
                    '''

    else:
        picard_summary, picard_histogram = outfiles

        statement = '''echo "Not compatible with SE data"
                       > %(picard_summary)s;
                       echo "Not compatible with SE data"
                       > %(picard_histogram)s
                    '''
    P.run(statement)


@merge(insertSizeMetricsAndHistograms,
       "qc.dir/qc_insert_size_metrics.load")
def loadInsertSizeMetrics(infiles, outfile):
    '''
    Load the insert size metrics to a single table of the project database.
    '''

    if PAIRED:
        picard_summaries = [x[0] for x in infiles]

        P.concatenate_and_load(picard_summaries, outfile,
                               regex_filename=(".*/.*/(.*)"
                                               ".insert.size.metrics.summary"),
                               cat="sample_id",
                               options='')

    else:
        statement = '''echo "Not compatible with SE data"
                       > %(outfile)s
                    '''
        P.run(statement)


@merge(insertSizeMetricsAndHistograms,
       "qc.dir/qc_insert_size_histogram.load")
def loadInsertSizeHistograms(infiles, outfile):
    '''
    Load the histograms to a single table of the project database.
    '''

    if PAIRED:
        picard_histograms = [x[1] for x in infiles]

        P.concatenate_and_load(
            picard_histograms, outfile,
            regex_filename=(".*/.*/(.*)"
                            ".insert.size.metrics.histogram"),
            cat="sample_id",
            options='-i "insert_size" -e')

    else:
        statement = '''echo "Not compatible with SE data"
                       > %(outfile)s
                    '''
        P.run(statement)


# -------------- No. reads mapping to spike-ins vs genome ------------------- #

@follows(mkdir("qc.dir/spike.vs.genome.dir"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           r"qc.dir/spike.vs.genome.dir/\1.uniq.mapped.reads")
def spikeVsGenome(infile, outfile):
    '''
    Summarise the number of reads mapping uniquely to spike-ins and genome.
    Compute the ratio of reads mapping to spike-ins vs genome.
    Only uniquely mapping reads are considered.'''

    # TODO: do this from featureCounts instead

    header = "\\t".join(["nreads_uniq_map_genome", "nreads_uniq_map_spike",
                        "fraction_spike"])

    statement = ''' echo -e "%(header)s" > %(outfile)s;
                    samtools view %(infile)s
                    | grep NH:i:1
                    | awk 'BEGIN{OFS="\\t";spikein=0;genome=0};
                           {total+=1};
                           $3~/%(spikein_pattern)s*/{spikein+=1};
                           END{frac=spikein/total;
                               genome=total-spikein;
                               print genome,spikein,frac};'
                    >> %(outfile)s
                '''
    P.run(statement)


@merge(spikeVsGenome,
       "qc.dir/qc_spike_vs_genome.load")
def loadSpikeVsGenome(infiles, outfile):
    '''
    Load number of reads uniquely mapping to genome & spike-ins
    and fraction of spike-ins to a single table of the project database.
    '''

    P.concatenate_and_load(infiles, outfile,
                           regex_filename=".*/.*/(.*).uniq.mapped.reads",
                           cat="sample_id",
                           options='-i "sample_id"')


# ------------------------- No. genes detected ------------------------------ #

@active_if(fastqMode)
@follows(mkdir("qc.dir/"), loadSalmonTPMs, loadEnsemblAnnotations)
@files("salmon.dir/salmon.genes.tpms.load",
       "qc.dir/number.genes.detected.salmon")
def numberGenesDetectedSalmon(infile, outfile):
    '''
    Count no genes detected at copynumer > 0 in each sample.
    '''

    table = P.to_table(infile)

    statement = '''select distinct s.*, i.gene_biotype
                   from %(table)s s
                   inner join transcript_info i
                   on s.gene_id=i.gene_id
                ''' % locals()

    df = DB.fetch_DataFrame(statement, DATABASE)

    melted_df = pd.melt(df, id_vars=["gene_id", "gene_biotype"])

    grouped_df = melted_df.groupby(["gene_biotype", "variable"])

    agg_df = grouped_df.agg({"value": lambda x:
                             np.sum([1 for y in x if y > 0])})
    agg_df.reset_index(inplace=True)

    count_df = pd.pivot_table(agg_df, index="variable",
                              values="value", columns="gene_biotype")
    count_df["total"] = count_df.apply(np.sum, 1)
    count_df["sample_id"] = count_df.index

    count_df.to_csv(outfile, index=False, sep="\t")


@active_if(fastqMode)
@follows(annotations)
@files(numberGenesDetectedSalmon,
       "qc.dir/qc_no_genes_salmon.load")
def loadNumberGenesDetectedSalmon(infile, outfile):
    '''
    Load the numbers of genes expressed to the project database.
    '''

    P.load(infile, outfile,
           options='-i "sample_id"')


@follows(annotations)
@files(loadFeatureCounts,
       "qc.dir/number.genes.detected.featurecounts")
def numberGenesDetectedFeatureCounts(infile, outfile):
    '''
    Count of genes detected by featureCount at counts > 0 in each sample.
    '''

    table = P.to_table(infile)

    # attach = '''attach "%(ANN_DATABASE)s" as anndb''' % globals()
    statement = '''select distinct h.*, gene_biotype
                   from %(table)s h
                   inner join transcript_info i
                   on h.gene_id=i.gene_id
               ''' % locals()

    melted_df = DB.fetch_DataFrame(statement, DATABASE)

    grouped_df = melted_df.groupby(["gene_biotype", "track"])

    agg_df = grouped_df.agg({"counts": lambda x:
                             np.sum([1 for y in x if y > 0])})
    agg_df.reset_index(inplace=True)

    count_df = pd.pivot_table(agg_df, index="track",
                              values="counts", columns="gene_biotype")
    count_df["total"] = count_df.apply(np.sum, 1)
    count_df["sample_id"] = count_df.index

    count_df.to_csv(outfile, index=False, sep="\t")


@files(numberGenesDetectedFeatureCounts,
       "qc.dir/qc_no_genes_featurecounts.load")
def loadNumberGenesDetectedFeatureCounts(infile, outfile):
    '''
    Load the numbers of genes expressed to the project database.
    '''

    P.load(infile, outfile,
           options='-i "sample_id"')


# --------------------- Fraction of spliced reads --------------------------- #

@follows(mkdir("qc.dir/fraction.spliced.dir/"))
@transform(collectBAMs,
           regex(r".*/(.*).bam"),
           r"qc.dir/fraction.spliced.dir/\1.fraction.spliced")
def fractionReadsSpliced(infile, outfile):
    '''
    Compute fraction of reads containing a splice junction.
    * paired-endedness is ignored
    * only uniquely mapping reads are considered.
    '''

    statement = '''echo "fraction_spliced" > %(outfile)s;
                   samtools view %(infile)s
                   | grep NH:i:1
                   | cut -f 6
                   | awk '{if(index($1,"N")==0){us+=1}
                           else{s+=1}}
                          END{print s/(us+s)}'
                   >> %(outfile)s
                 '''

    P.run(statement)


@merge(fractionReadsSpliced,
       "qc.dir/qc_fraction_spliced.load")
def loadFractionReadsSpliced(infiles, outfile):
    '''
    Load fractions of spliced reads to a single table of the project database.
    '''

    P.concatenate_and_load(infiles, outfile,
                           regex_filename=".*/.*/(.*).fraction.spliced",
                           cat="sample_id",
                           options='-i "sample_id"')


# ---------------- Prepare a post-mapping QC summary ------------------------ #

@follows(mkdir("annotations.dir"))
@files(None,
       "annotations.dir/sample.information.txt")
def sampleInformation(infiles, outfile):
    '''
    Make a database table containing per-sample information.
    '''

    SAMPLES.to_csv(outfile, index=False, sep="\t")


@transform(sampleInformation,
           suffix(".txt"),
           ".load")
def loadSampleInformation(infile, outfile):
    '''
    Load the sample information table to the project database.
    '''

    P.load(infile, outfile)


@merge([loadSampleInformation,
        loadCollectRnaSeqMetrics,
        loadThreePrimeBias,
        loadEstimateLibraryComplexity,
        loadSpikeVsGenome,
        loadFractionReadsSpliced,
        loadNumberGenesDetectedSalmon,
        loadNumberGenesDetectedFeatureCounts,
        loadAlignmentSummaryMetrics,
        loadInsertSizeMetrics],
       "qc.dir/qc_summary.txt")
def qcSummary(infiles, outfile):
    '''
    Create a summary table of relevant QC metrics.
    '''

    # Some QC metrics are specific to paired end data
    if PAIRED:
        exclude = []
        paired_columns = '''READ_PAIRS_EXAMINED as no_pairs,
                              PERCENT_DUPLICATION as pct_duplication,
                              ESTIMATED_LIBRARY_SIZE as library_size,
                              PCT_READS_ALIGNED_IN_PAIRS
                                       as pct_reads_aligned_in_pairs,
                              MEDIAN_INSERT_SIZE
                                       as median_insert_size,
                           '''
        pcat = "PAIR"

    else:
        exclude = ["qc_library_complexity", "qc_insert_size_metrics"]
        paired_columns = ''
        pcat = "UNPAIRED"

    if fastqMode:
        exclude = exclude
        fastq_columns = '''qc_no_genes_salmon.protein_coding
                              as salmon_no_genes_pc,
                           qc_no_genes_salmon.total
                              as salmon_no_genes,
                        '''
    else:
        exclude = exclude + ["qc_no_genes_salmon"]
        fastq_columns = ''

    tables = [P.to_table(x) for x in infiles
              if P.to_table(x) not in exclude]

    t1 = tables[0]

    name_fields = PARAMS["name_field_titles"].strip()

    stat_start = '''select distinct %(name_fields)s,
                                    sample_information.sample_id,
                                    fraction_spliced,
                                    fraction_spike,
                                    %(fastq_columns)s
                                    qc_no_genes_featurecounts.protein_coding
                                       as featurecounts_no_genes_pc,
                                    qc_no_genes_featurecounts.total
                                       as featurecounts_no_genes,
                                    three_prime_bias
                                       as three_prime_bias,
                                    nreads_uniq_map_genome,
                                    nreads_uniq_map_spike,
                                    %(paired_columns)s
                                    PCT_MRNA_BASES
                                       as pct_mrna,
                                    PCT_CODING_BASES
                                       as pct_coding,
                                    PCT_PF_READS_ALIGNED
                                       as pct_reads_aligned,
                                    TOTAL_READS
                                       as total_reads,
                                    PCT_ADAPTER
                                       as pct_adapter,
                                    PF_HQ_ALIGNED_READS*1.0/PF_READS
                                       as pct_pf_reads_aligned_hq
                   from %(t1)s
                ''' % locals()

    join_stat = ""
    for table in tables[1:]:
        join_stat += "left join " + table + "\n"
        join_stat += "on " + t1 + ".sample_id=" + table + ".sample_id\n"

    where_stat = '''where qc_alignment_summary_metrics.CATEGORY="%(pcat)s"
                 ''' % locals()

    statement = "\n".join([stat_start, join_stat, where_stat])

    df = DB.fetch_DataFrame(statement, PARAMS["database_file"])
    df.to_csv(outfile, sep="\t", index=False)


@transform(qcSummary,
           suffix(".txt"),
           ".load")
def loadQCSummary(infile, outfile):
    '''
    Load summary to project database.
    '''

    P.load(infile, outfile)


@follows(loadQCSummary, loadInsertSizeHistograms)
def qc():
    '''
    Target for executing quality control.
    '''
    pass


# --------------------- < generic pipeline tasks > -------------------------- #

@follows(mkdir("notebook.dir"))
@transform(glob.glob(os.path.join(os.path.dirname(__file__),
                                  "pipeline_notebooks",
                                  os.path.basename(__file__)[:-len(".py")],
                                  "*")),
           regex(r".*/(.*)"),
           r"notebook.dir/\1")
def notebooks(infile, outfile):
    '''
    Utility function to copy the notebooks from the source directory
    to the working directory.
    '''

    shutil.copy(infile, outfile)


@follows(quantitation, qc)
def full():
    pass


print(sys.argv)

if __name__ == "__main__":
    sys.exit(P.main(sys.argv))
