#!/usr/bin/env python

""" MultiQC submodule to parse output from Picard VariantCallingMetrics """

import logging
from collections import OrderedDict
from multiqc.plots import bargraph

# Initialise the logger
log = logging.getLogger(__name__)


def parse_reports(parent_module):
    """ Find Picard VariantCallingMetrics reports and process their data """

    # get data
    data = collect_data(parent_module)

    # Filter to strip out ignored sample names
    data = parent_module.ignore_samples(data)

    # Reference data in parent module
    parent_module.picard_variantCalling_data = data

    if len(data) > 0:
        derive_data(data)

        # Write parsed data to a file
        parent_module.write_data_file(data, 'multiqc_picard_variantCalling')

        parent_module.general_stats_headers['DBSNP_TITV'] = {
            'title': 'TiTV ratio (known)',
            'description': "The Transition/Transversion ratio of the passing bi-allelic SNP calls made at SNP-database sites.",
            'min': 0,
            'scale': 'Blues',
            'shared_key': 'titv_ratio'
        }

        parent_module.general_stats_headers['NOVEL_TITV'] = {
            'title': 'TiTV ratio (novel)',
            'description': "The Transition/Transversion ratio of the passing bi-allelic SNP calls made at non-SNP-database sites.",
            'min': 0,
            'scale': 'Blues',
            'shared_key': 'titv_ratio'
        }

        parent_module.general_stats_headers['DBSNP_INS_DEL_RATIO'] = {
            'title': 'InDel ratio (known)',
            'description': "The Insertion / Deletion ratio of the passing bi-allelic SNP calls made at SNP-database sites.",
            'min': 0,
            'scale': 'Greens',
            'shared_key': 'indel_ratio',
            'hidden': True,
        }

        parent_module.general_stats_headers['NOVEL_INS_DEL_RATIO'] = {
            'title': 'InDel ratio (novel)',
            'description': "The Insertion / Deletion ratio of the passing bi-allelic SNP calls made at non-SNP-database sites.",
            'min': 0,
            'scale': 'Greens',
            'shared_key': 'indel_ratio',
            'hidden': True,
        }


        parent_module.general_stats_headers['total_called_variants_known'] = {
            'title': 'Called Variants (known)',
            'description': "Total counts of variants in SNP-database sites.",
            'shared_key': 'variant_count',
            'min': 0,
            'format': '{0:,.0f}',
            'hidden': True,
        }

        parent_module.general_stats_headers['total_called_variants_novel'] = {
            'title': 'Called Variants (novel)',
            'description': "Total counts of variants in non-SNP-database sites.",
            'shared_key': 'variant_count',
            'min': 0,
            'format': '{0:,.0f}',
            'hidden': True,
        }

        for s_name in data:
            if s_name not in parent_module.general_stats_data:
                parent_module.general_stats_data[s_name] = dict()
            parent_module.general_stats_data[s_name].update(data[s_name])


        # Variant Counts Bargraph
        parent_module.add_section(
            name='Variant Types',
            anchor='picard-variants-types',
            description='Variants that have been called, looking at variant types. Optionally filtered on label.',
            helptext="""
            Only passing variants are shown (i.e. non-filtered).\n
            SNPs are bi-allelic.\n
            Complex InDels are both an insertion and a deletion.
            """,
            plot=compare_variant_type_plot(data)
        )

        # Variant Counts Table
        parent_module.add_section(
            name='Variant Labels',
            anchor='picard-variants-labels',
            description='Variants that have been called, comparing with known variant sites.',
            helptext="""
            Only passing variants are shown (i.e. non-filtered).\n
            Variants contain bi-allelic SNPs, multi-allelic SNPs, simple and complex inserts and deletions.
            """,
            plot=compare_variants_label_plot(data)
        )

    return len(data)


def collect_data(parent_module):
    """ Find Picard VariantCallingMetrics reports and parse their data """

    data = dict()
    for file_meta in parent_module.find_log_files('picard/variant_calling_metrics', filehandles=True):
        s_name = None
        for header, value in table_in(file_meta['f'], pre_header_string='## METRICS CLASS'):
            if header == 'SAMPLE_ALIAS':
                s_name = value
                if s_name in data:
                    log.debug("Duplicate sample name found in {}! Overwriting: {}".format(file_meta['fn'], s_name))
                data[s_name] = OrderedDict()
            else:
                data[s_name][header] = value
    return data


def table_in(filehandle, pre_header_string):
    """ Generator that assumes a table starts the line after a given string """

    in_histogram = False
    next_is_header = False
    headers = list()
    for line in stripped(filehandle):
        if not in_histogram and line.startswith(pre_header_string):
            in_histogram = True
            next_is_header = True
        elif in_histogram and next_is_header:
            next_is_header = False
            headers = line.split("\t")
        elif in_histogram:
            values = line.split("\t")
            if values != ['']:
                for couple in zip(headers, values):
                    yield couple


def derive_data(data):
    """ Based on the data derive additional data """

    for s_name, values in data.items():
        # setup holding variable

        # Sum all variants that have been called
        total_called_variants = 0
        for value_name in ['TOTAL_SNPS', 'TOTAL_COMPLEX_INDELS', 'TOTAL_MULTIALLELIC_SNPS', 'TOTAL_INDELS']:
            total_called_variants = total_called_variants + int(values[value_name])
        values['total_called_variants'] = total_called_variants

        # Sum all variants that have been called and are known
        total_called_variants_known = 0
        for value_name in ['NUM_IN_DB_SNP', 'NUM_IN_DB_SNP_COMPLEX_INDELS', 'NUM_IN_DB_SNP_MULTIALLELIC']:
            total_called_variants_known = total_called_variants_known + int(values[value_name])
        total_called_variants_known = total_called_variants_known + int(values['TOTAL_INDELS']) - int(values['NOVEL_INDELS'])
        values['total_called_variants_known'] = total_called_variants_known

        # Extrapolate the total novel variants
        values['total_called_variants_novel'] = total_called_variants - total_called_variants_known


def stripped(iterator):
    """ Generator to strip string of whitespace """
    for item in iterator:
        yield item.strip()


def compare_variant_type_plot(data):
    """ Return HTML for the Variant Counts barplot """
    keys = OrderedDict()
    keys['snps'] = {'name': 'SNPs', 'color': '#7cb5ec'}
    keys['indels'] = {'name': 'InDels', 'color': '#90ed7d'}
    keys['multiallelic_snps'] = {'name': 'multi-allelic SNP', 'color': 'orange'}
    keys['complex_indels'] = {'name': 'Complex InDels', 'color': '#8085e9'}

    total_variants = dict()
    known_variants = dict()
    novel_variants = dict()
    for s_name, values in data.items():
        total_variants[s_name] = {
            'snps': values['TOTAL_SNPS'],
            'indels': values['TOTAL_INDELS'],
            'multiallelic_snps': values['TOTAL_MULTIALLELIC_SNPS'],
            'complex_indels': values['TOTAL_COMPLEX_INDELS'],
        }

        known_variants[s_name] = {
            'snps': values['NUM_IN_DB_SNP'],
            'indels': int(values['TOTAL_INDELS'])-int(values['NOVEL_INDELS']),
            'multiallelic_snps': values['NUM_IN_DB_SNP_MULTIALLELIC'],
            'complex_indels': values['NUM_IN_DB_SNP_COMPLEX_INDELS'],
        }

        novel_variants[s_name] = {
            'snps': values['NOVEL_SNPS'],
            'indels': values['NOVEL_INDELS'],
            'multiallelic_snps': int(values['TOTAL_MULTIALLELIC_SNPS'])-int(values['NUM_IN_DB_SNP_MULTIALLELIC']),
            'complex_indels': int(values['TOTAL_COMPLEX_INDELS'])-int(values['NUM_IN_DB_SNP_COMPLEX_INDELS']),
        }

    plot_conf = {
        'id': 'picard_variantCallingMetrics_variant_type',
        'title': 'Picard: Variants Called',
        'ylab': 'Counts of Variants',
        'hide_zero_cats': False,
        'data_labels': [
            {'name': 'Total'},
            {'name': 'Known'},
            {'name': 'Novel'}
        ],
    }
    return bargraph.plot(data=[total_variants, known_variants, novel_variants],
                         cats=[keys, keys, keys],
                         pconfig=plot_conf)


def compare_variants_label_plot(data):
    """ Return HTML for the Compare variants plot"""
    keys = OrderedDict()

    keys['total_called_variants_known'] = {'name': 'Known Variants'}
    keys['total_called_variants_novel'] = {'name': 'Novel Variants'}

    pconfig = {
        'id': 'picard_variantCallingMetrics_variant_label',
        'title': 'Picard: Variants Called',
        'ylab': 'Counts of Variants',
     }

    return bargraph.plot(data, cats=keys, pconfig=pconfig)


def merge_two_dicts(x, y):
    z = x.copy()   # start with x's keys and values
    z.update(y)    # modifies z with y's keys and values & returns None
    return z
