"""
run the pizzly fusion caller for RNA-seq
https://github.com/pmelsted/pizzly
http://www.biorxiv.org/content/early/2017/07/20/166322
"""
from __future__ import print_function

import os

from bcbio import utils
import bcbio.pipeline.datadict as dd
from bcbio.pipeline import config_utils
from bcbio.distributed.transaction import file_transaction
from bcbio.rnaseq import kallisto, sailfish, gtf
from bcbio.provenance import do
from bcbio.utils import file_exists, safe_makedir

h5py = utils.LazyImport("h5py")
import numpy as np

def get_fragment_length(data):
    """
    lifted from
    https://github.com/pmelsted/pizzly/scripts/pizzly_get_fragment_length.py
    """
    h5 = kallisto.get_kallisto_h5(data)
    cutoff = 0.95
    with h5py.File(h5) as f:
        x = np.asarray(f['aux']['fld'], dtype='float64')
    y = np.cumsum(x)/np.sum(x)
    fraglen = np.argmax(y > cutoff)
    return(fraglen)

def run_pizzly(data):
    work_dir = dd.get_work_dir(data)
    pizzlydir = os.path.join(work_dir, "pizzly")
    samplename = dd.get_sample_name(data)
    gtf = dd.get_gtf_file(data)
    if dd.get_transcriptome_fasta(data):
        gtf_fa = dd.get_transcriptome_fasta(data)
    else:
        gtf_fa = sailfish.create_combined_fasta(data)
    fraglength = get_fragment_length(data)
    cachefile = os.path.join(pizzlydir, "pizzly.cache")
    fusions = kallisto.get_kallisto_fusions(data)
    pizzlypath = config_utils.get_program("pizzly", dd.get_config(data))
    outdir = pizzly(pizzlypath, gtf, gtf_fa, fraglength, cachefile, pizzlydir,
                    fusions, samplename, data)
    return outdir

def pizzly(pizzly_path, gtf, gtf_fa, fraglength, cachefile, pizzlydir, fusions,
           samplename, data):
    outdir = os.path.join(pizzlydir, samplename)
    pizzly_gtf = make_pizzly_gtf(gtf, os.path.join(pizzlydir, "pizzly.gtf"), data)
    with file_transaction(data, outdir) as tx_out_dir:
        safe_makedir(tx_out_dir)
        out_stem = os.path.join(tx_out_dir, samplename)
        cmd = ("{pizzly_path} -k 31 --gtf {pizzly_gtf} --cache {cachefile} "
            "--align-score 2 --insert-size {fraglength} --fasta {gtf_fa} "
            "--output {out_stem} {fusions}")
        message = ("Running pizzly on %s." % fusions)
        do.run(cmd.format(**locals()), message)
        pizzlycalls = out_stem + ".json"
        flatfile = out_stem + "-flat.tsv"
        flatten_pizzly(pizzlycalls, flatfile, data)
    return outdir

def make_pizzly_gtf(gtf_file, out_file, data):
    """
    pizzly needs the GTF to be in gene -> transcript -> exon order for each
    gene. it also wants the gene biotype set as the source
    """
    if file_exists(out_file):
        return out_file
    db = gtf.get_gtf_db(gtf_file)
    with file_transaction(data, out_file) as tx_out_file:
        with open(tx_out_file, "w") as out_handle:
            for gene in db.features_of_type("gene"):
                children = [x for x in db.children(id=gene)]
                for child in children:
                    if child.attributes.get("gene_biotype", None):
                        gene_biotype = child.attributes.get("gene_biotype")
                gene.attributes['gene_biotype'] = gene_biotype
                gene.source = gene_biotype[0]
                print(gene, file=out_handle)
                for child in children:
                    child.source = gene_biotype[0]
                    # gffread produces a version-less FASTA file
                    child.attributes.pop("transcript_version", None)
                    print(child, file=out_handle)
    return out_file

def flatten_pizzly(in_file, out_file, data):
    pizzlyflatten = config_utils.get_program("pizzly_flatten_json.py", data)
    if file_exists(out_file):
        return out_file
    cmd = "{pizzlyflatten} {in_file} > {tx_out_file}"
    message = "Flattening {in_file} to {out_file}."
    with file_transaction(data, out_file) as tx_out_file:
        do.run(cmd.format(**locals()), message.format(**locals()))
    return out_file

@@ -524,16 +524,36 @@
 py as np
+%0Aimport pandas as pd
 %0A%0Adef ge
@@ -2403,39 +2403,146 @@
    f
-latten_pizzly(pizzlycalls, flat
+ilteredfile = out_stem + %22-flat-filtered.tsv%22%0A        flatten_pizzly(pizzlycalls, flatfile, data)%0A        filter_pizzly(flatfile, filtered
 file
@@ -4115,28 +4115,357 @@
 als()))%0A    return out_file%0A
+%0Adef filter_pizzly(in_file, out_file, data):%0A    df = pd.read_csv(in_file, header=0, sep=%22%5Ct%22)%0A    df = df.query('paircount %3E 1 and splitcount %3E 1')%0A    if file_exists(out_file):%0A        return out_file%0A    with file_transaction(out_file) as tx_out_file:%0A        df.to_csv(tx_out_file, sep=%22%5Ct%22, index=False)%0A    return out_file%0A

Add filtering of low evidence pizzly fusions.