#!/usr/bin/env python3
"""
This module is used for parsing generated fusion files from fusion tools and generating supporting
files needed for other tools like: create_mqc_section.py and generate_report.py
"""
import argparse
import os.path
import sys
import json
from yaml import dump

TOOLS = ['star_fusion', 'fusioncatcher', 'ericscript', 'pizzly', 'squid']
SUMMARY = 'summary.yaml'
OUTPUT = 'fusions.txt'

def save(p_output, p_fusions):
    """
    Makes union of detected fusions and previous fusions found by other tools. Saves only unique
    fusions into OUTPUT.
    Args:
        p_output (string): OUTPUT (fusions.txt)
        p_fusions (string): concatenated list of all found fusions by a specific tool
            Example: BRAF--AKAP9\nBRD4_NUTM1
    """
    old = p_output.read().split('\n')
    new = p_fusions.split('\n')
    unique_fusions = set().union(old, new)
    unique_fusions = '\n'.join(unique_fusions).lstrip()
    # clear file
    p_output.seek(0)
    p_output.truncate()
    # write fusions
    p_output.write(unique_fusions)

def fi_format(p_gene1, p_gene2):
    """Set predefined format FUSION1--FUSION2
    Args:
        p_gene1 (string): left gene of the fusion
        p_gene2 (string): right gene of the fusion
    """
    return '{}--{}\n'.format(p_gene1, p_gene2)

def squid(p_file):
    """Extracting fusions from Squid output file
    Args:
        p_file (FILE): Output file generated by a specific fusion tool
    """
    fusions = ''
    next(p_file)    # skip header
    for line in p_file:
        tmp = line.rstrip('\n').split('\t')[11].split(':')
        if len(tmp) == 2:
            fusions += fi_format(tmp[0], tmp[1])

    return fusions

def pizzly(p_file):
    """Extracting fusions from Pizzly output file
    Args:
        p_file (FILE): Output file generated by a specific fusion tool
    """
    fusions = ''
    data = json.load(p_file)
    for item in data['genes']:
        fusions += fi_format(item['geneA']['name'], item['geneB']['name'])

    return fusions

def ericscript(p_file):
    """Extracting fusions from Ericscript output file
    Args:
        p_file (FILE): Output file generated by a specific fusion tool
    """
    fusions = ''
    next(p_file)    # skip header
    for line in p_file:
        tmp = line.split('\t')
        fusions += fi_format(tmp[0], tmp[1])

    return fusions

def star_fusion(p_file):
    """Extracting fusions from STAR-Fusion output file
    Args:
        p_file (FILE): Output file generated by a specific fusion tool
    """
    fusions = ''
    next(p_file)    #skip header
    for line in p_file:
        tmp = line.split('\t')[0].split('--')
        fusions += fi_format(tmp[0], tmp[1])

    return fusions

def fusioncatcher(p_file):
    """Extracting fusions from Fusioncatcher output file
    Args:
        p_file (FILE): Output file generated by a specific fusion tool
    """
    fusions = ''
    next(p_file)    # skip header
    for line in p_file:
        tmp = line.split('\t')
        fusions += fi_format(tmp[0], tmp[1])

    return fusions

def transform(p_input, p_tool, p_output):
    """Dynamically executes function for fusion extraction (functions are named based on the tools).
    It generates 2 files:
        SUMMARY: list of tools names and corresponding found fusions
        OUTPUT: list of all uniques fusions found by executed tools
    Args:
        p_input (string): Output file generated by a specific fusion tool
        p_tool (string): name of the tool defined from command line
        p_output (string): OUTPUT
    """
    errors = []
    if not os.path.exists(p_input):
        errors.append('Defined {} doesn\'t exist'.format(p_input))

    if os.stat(p_input).st_size == 0:
        errors.append('Provided file {} is empty'.format(p_input))

    if p_tool not in TOOLS:
        errors.append('Defined {} not in the supported list of transformations!'.format(p_tool))

    if errors:
        map(print, errors)
    else:
        if not os.path.exists(OUTPUT):
            with open(OUTPUT, 'w'):
                pass

        try:
            fusions = ''
            with open(p_input, 'r') as in_file:
                func = getattr(sys.modules[__name__], p_tool)   # get function from parameter
                fusions = func(in_file).rstrip()  # call function
                in_file.close()

            with open(p_output, 'r+') as out_file, open(SUMMARY, 'a') as summary:
                if fusions:
                    save(out_file, fusions)
                    summary.write(
                        dump(
                            {p_tool : list(set(fusions.split('\n')))},
                            default_flow_style=False,
                            allow_unicode=True)
                    )
                else:
                    summary.write(
                        dump(
                            {p_tool: None},
                            default_flow_style=False,
                            allow_unicode=True
                        )
                    )
                out_file.close()
                summary.close()
        except IOError as error:
            print(error)

def main():
    """Main function for processing command line arguments"""
    parser = argparse.ArgumentParser(
        description='Tool for extracting fusion-genes from various fusion detection tools'
    )
    parser.add_argument(
        '-i', '--input',
        help='Output file of a fusion tool',
        type=str,
        required=True
    )
    parser.add_argument(
        '-t', '--tool',
        help='',
        choices='|'.join(TOOLS),
        type=str,
        required=True
    )
    args = parser.parse_args()
    transform(args.input, args.tool, OUTPUT)

if __name__ == "__main__":
    main()
