#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
charm-cli.py: Simple command line interface for CHarm.
"""

import argparse
import logging

try:
    import matplotlib

    matplotlib.use('Agg')
    matplotlib.rc('font', **{'sans-serif': 'DejaVu Sans',
                             'serif': 'DejaVu Serif',
                             'family': 'sans-serif'})
    import matplotlib.pyplot
except ImportError as e:
    print('ERROR: {}'.format(e.msg))
    exit(1)
try:
    import numpy
except ImportError as e:
    print('ERROR: {}'.format(e.msg))
    exit(1)

try:
    from LibCharm.Sequence import Sequence
    from LibCharm import IO
except ImportError as e:
    print('ERROR: {}'.format(e.msg))
    exit(1)


def autolabel(rects, ax, labels, vertical=True):
    """
    Automatically adds labels above a bar in a bar graph.

    :param rects:    list of bars to be labelled (e.g. generated by ax.bar())
    :param ax:       axis (axis object from matplotlib)
    :param labels:   list of labels
    :param vertical: rotate the labels by 90° if true

    """
    if vertical:
        rotation = 'vertical'
    else:
        rotation = 'horizontal'

    if len(labels) == len(rects):
        heights = []
        for rect in rects:
            height = rect.get_height()
            heights.append(height)

        max_height = max(heights)

        for rect in rects:
            i = rects.index(rect)
            label = labels[i]
            height = rect.get_height()
            if height > 0:
                y = 1.05 * height
            else:
                y = 0.02 * max_height
            ax.text(rect.get_x() + rect.get_width() / 2., y, str(label),
                    ha='center', va='bottom', rotation=rotation, size='x-small')


def plot_codon_usage(sequence, ax):
    """
    Plot the codon usage for origin and target host as bar graph

    :param sequence: LibCharm.Sequence object
    :param ax      : matplotlib axis object
    """

    x1 = x2 = numpy.arange(len(sequence.codons))
    bar_width = 0.5
    xlabels = []

    origin_f = []
    target_f = []

    # extract data to plot from sequence object
    for c in sequence.codons:
        origin_f.append(c['origin_f'])
        target_f.append(c['target_f'])
        xlabels.append(c['aa'])

    # convert lists to numpy arrays
    origin_f = numpy.array(origin_f)
    target_f = numpy.array(target_f)

    # plot data
    p1 = ax.bar(x1, origin_f, color='b', width=bar_width)
    p2 = ax.bar(x2 + (0.5 * bar_width), target_f, color='r', width=bar_width)

    # hide top and right axes
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    # set tick parameters
    ax.tick_params(axis='both', which='both', direction='out')
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()
    # position xticks and labels on x axis to be centered for both bars
    ax.set_xticks(x1 + bar_width / 2)
    ax.set_xticklabels(xlabels, **{'family': 'monospace'})
    ax.set_xlabel('amino acid')
    # add a legend to the plot
    ax.legend((p1, p2), ('Origin organism', 'Host organism'), loc=2, bbox_to_anchor=(1, 1))
    ax.hlines(sequence.lower_threshold, 0, len(x1), colors='k', linestyles='solid', **{'linewidth': 1})

    if not sequence.use_frequency:
        # set the y axis label
        ax.set_ylabel('codon usage [fraction]')
        # specify the distance between the ticks on the y axis
        major_locator = matplotlib.ticker.MultipleLocator(0.1)
        minor_locator = matplotlib.ticker.MultipleLocator(0.01)
    else:
        # set the y axis label if frequency is used instead of fractions
        ax.set_ylabel('codon usage [frequency/1000]')
        # specify the distance between the ticks on the y axis
        major_locator = matplotlib.ticker.MultipleLocator(10)
        minor_locator = matplotlib.ticker.MultipleLocator(1)

    # set the distance between the ticks on the y axis
    ax.yaxis.set_major_locator(major_locator)
    ax.yaxis.set_minor_locator(minor_locator)


def plot_codon_usage_differences(sequence, ax):
    """
    Plot the difference in codon usage for origin and target host as bar graph

    :param sequence: LibCharm.Sequence object
    :param ax:       matplotlib axis object
    """

    # Generate a range of residues out of the length of the sequence array
    x1 = numpy.arange(len(sequence.codons))

    # Set the threshold according to use_frequency
    if sequence.use_frequency:
        threshold = 5
    else:
        threshold = 0.2

    # Set width of bars
    bar_width = 0.8
    # Initialize array of labels for the x axis
    xlabels = []

    # Initialize arrays of data and labels for the bars
    df = []
    bar_labels = []

    # walk over the codons in sequence
    for c in sequence.codons:
        # add final_df to data array
        df.append(c['final_df'])
        # add residue to xlabels
        xlabels.append(c['aa'])
        # generate bar label and add to list
        label = u'{} → {}'.format(c['original'], c['new'])
        bar_labels.append(label)

    # convert lists to numpy arrays
    bar_labels = numpy.array(bar_labels)
    df = numpy.array(df)
    # find bars that exceed the threshold
    mask1 = numpy.ma.where(df > threshold)
    mask2 = numpy.ma.where(df <= threshold)

    # plot and color bars accordingly
    p1 = ax.bar(x1[mask1], df[mask1], color='r', width=bar_width)
    autolabel(p1, ax, bar_labels[mask1], vertical=True)

    p2 = ax.bar(x1[mask2], df[mask2], color='b', width=bar_width)
    autolabel(p2, ax, bar_labels[mask2], vertical=True)

    # hide top and right axis
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.tick_params(axis='both', which='both', direction='out')
    ax.get_xaxis().tick_bottom()
    ax.get_yaxis().tick_left()

    # set x axis labels to be centered and to use a monospaced font
    ax.set_xticks(x1 + bar_width / 2)
    ax.set_xticklabels(xlabels, **{'family': 'monospace'})
    ax.set_xlabel('amino acid')

    ax.set_ylabel(r'Differential codon usage $f_{origin} - f_{host}$')

    if not sequence.use_frequency:
        major_locator = matplotlib.ticker.MultipleLocator(0.05)
        minor_locator = matplotlib.ticker.MultipleLocator(0.01)
    else:
        major_locator = matplotlib.ticker.MultipleLocator(10)
        minor_locator = matplotlib.ticker.MultipleLocator(1)

    ax.legend((p1, p2), (u'Δf > {}'.format(threshold), u'Δf ≤ {}'.format(threshold)), loc=2, bbox_to_anchor=(1, 1))

    ax.yaxis.set_major_locator(major_locator)
    ax.yaxis.set_minor_locator(minor_locator)

    ax.hlines(threshold, 0, len(x1), colors='k', linestyles='dotted', **{'linewidth': 1})


def plot(sequence, prefix=None):
    """
    Wrapper for plot_codon_usage_differences and plot_codon_usage

    :param sequence:  LibCharm.Sequence object
    :param prefix:    Resulting plot files will be prefixed with 'prefix'
    """

    if prefix:
        filename = '{}_charm_results.svg'.format(prefix)
    else:
        filename = 'charm_results.svg'

    # Create a plot with two subplots
    fig, axarr = matplotlib.pyplot.subplots(2, figsize=(50, 20), dpi=300)

    # Actually plot data
    plot_codon_usage(sequence, axarr[0])
    plot_codon_usage_differences(sequence, axarr[1])

    # Save plot as svg
    matplotlib.pyplot.savefig(filename, format='svg', orientation='landscape', papertype='a4')


def parse_arguments():
    """
    Parse command line arguments and return list of arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', action='store_true', help='increase output verbosity')
    parser.add_argument('-p', '--prefix', type=str, help='prefix for output files')
    parser.add_argument('-f', '--frequency', action='store_true', help='use frequency/1000 instead of fraction')
    parser.add_argument('-l', '--lower_frequency_alternative', action='store_true',
                        help='if two codons result in the same difference in codon usage '
                             'between origin and target host, use the lower frequency alternative')
    parser.add_argument('-t', '--threshold', type=float,
                        help='Lower threshold of codon usage. Defaults to 0.1 and 5 for fraction and '
                             'frequency respectively')
    parser.add_argument('-to', '--translation_table_origin', type=int,
                        help='id of translation table; Default is: standard genetic code = 1; '
                             'id corresponds to \'trans_table\' '
                             'on http://www.ncbi.nlm.nih.gov/Taxonomy/Utils/wprintgc.cgi')
    parser.add_argument('-th', '--translation_table_host', type=int,
                        help='id of translation table; Default is: standard genetic code = 1; '
                             'id corresponds to \'trans_table\' '
                             'on http://www.ncbi.nlm.nih.gov/Taxonomy/Utils/wprintgc.cgi')
    parser.add_argument('origin', type=int, help='species id of origin organism taken from '
                                                 '\'http://www.kazusa.or.jp/codon\' (e.g. \'83333\' for E. coli K12)')
    parser.add_argument('host', type=int, help='species id of host organism taken from '
                                               '\'http://www.kazusa.or.jp/codon\' (e.g. \'83333\' for E. coli K12)')
    parser.add_argument('input', type=str, help='input file in FASTA format')
    args = parser.parse_args()

    return args


def initialize_logger(prefix):
    """
    Initialization of logging subsystem. Two logging handlers are brought up:
    'fh' which logs to a log file and 'ch' which logs to standard output.

    :param prefix:  prefix that is added to the filename
    :return logger: return a logger instance
    """
    logger = logging.getLogger('charm-cli')
    logger.setLevel(logging.INFO)

    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    logger.addHandler(ch)

    try:
        if prefix:
            log_filename = '{}_charm-cli.log'.format(prefix)
        else:
            log_filename = 'charm-cli.log'
        fh = logging.FileHandler(log_filename, 'w')
        fh.setLevel(logging.INFO)
        logger.addHandler(fh)
    except IOError as error:
        logger.warning('WARNING: Cannot create log file! Run charm-cli from a directory to '
                       'which you have write access.')
        logger.warning(error.msg)
        pass

    return logger


def main():
    """
    Main function of charm-cli.py.
    """

    # Parse command line arguments
    args = parse_arguments()
    # Initialize logging
    logger = initialize_logger(args.prefix)

    # Set translation tables according to user input. Defaults to standard genetic code (table 1)
    if args.translation_table_origin:
        translation_table_origin = args.translation_table_origin
    else:
        translation_table_origin = 1

    if args.translation_table_host:
        translation_table_host = args.translation_table_host
    else:
        translation_table_host = 1

    # set threshold if provided by the user and otherwise fall back to defaults
    if args.threshold:
        lower_threshold = args.threshold
    elif args.frequency:
        lower_threshold = 5
    else:
        lower_threshold = 0.1

    # initialize Sequence object with user provided input
    sequence = Sequence(IO.load_file(args.input), args.origin, args.host,
                        translation_table_origin=translation_table_origin,
                        translation_table_host=translation_table_host,
                        use_frequency=args.frequency,
                        lower_threshold=lower_threshold,
                        lower_alternative=args.lower_frequency_alternative)

    # harmonize the provided sequence
    harmonized_codons = sequence.get_harmonized_codons()
    # check if input and output sequence are identical
    verify_sequence = sequence.verify_harmonized_sequence()

    # log summary to standard output and log file
    logger.info('SUMMARY:\n')
    if verify_sequence:
        text = 'Success! Translation of harmonized and original sequence match:\n\n' \
               '{}\n'.format(sequence.harmonized_translated_sequence)
        logger.info(text)
    else:
        logger.error('ERROR: Translations of harmonized and original sequence DO NOT match!')
    logger.info('Harmonized codons: {}\n'.format(len(harmonized_codons)))

    df_above_thresh = 0
    for c in sequence.codons:
        if c['final_df'] > 0.2:
            df_above_thresh += 1

    if df_above_thresh > 0:
        logger.warning("WARNING: Difference in origin and target host codon usage of {} out of {} codons ({}%) exceeds 20%!\n".format(df_above_thresh,
                                                                                                                                      len(sequence.codons),
                                                                                                                                      round(df_above_thresh/len(sequence.codons)*100, 1)))
    else:
        logger.info("Differences of codon usage in origin and target host are within 20%.\n")

    table_header = '{:<10} {:^3} {:^4}    {:^4} {:^7} {:>6} {:<7} {:>6}'.format('position', 'aa', 'orig', 'new',
                                                                                'initial', 'final', 'origin', 'target')
    logger.info(table_header)

    warnings = []

    # Iterate over all codons in the sequence and print some statistics and information
    for c in sequence.codons:
        if str(c['original']) != str(c['new']):
            line = '{:<10} {:^3} {:<4} -> {:<4} {:<5.2f} -> {:<3.2f}  {:<5.2f} -> {:<3.2f}'.format(c['position'],
                                                                                                   c['aa'],
                                                                                                   c['original'],
                                                                                                   c['new'],
                                                                                                   c['initial_df'],
                                                                                                   c['final_df'],
                                                                                                   c['origin_f'],
                                                                                                   c['target_f'])
        else:
            line = '{:<10} {:^3} {:<12} {:<5.2f}          {:<5.2f} -> {:<3.2f}'.format(c['position'],
                                                                                       c['aa'],
                                                                                       c['original'],
                                                                                       c['initial_df'],
                                                                                       c['origin_f'],
                                                                                       c['target_f'])
        if c['ambiguous']:
            line += ' WARNING: Original codon is ambiguous!'
            warnings.append('Codon {} ({}) coding for {} is ambiguous! {} was chosen for the '
                            'harmonized sequence!'.format(c['position'],
                                                          c['original'],
                                                          c['aa'],
                                                          c['new']))

        logger.info(line)

    logger.info('\nCodon-harmonized sequence:\n\n{}'.format(sequence.harmonized_sequence))
    if warnings:
        logger.warn('\nWARNINGS OCCURRED DURING HARMONIZATION:\n')
        for warning in warnings:
            logger.warn(warning)

    plot(sequence, args.prefix)

    # Exit gracefully
    exit(0)


if __name__ == "__main__":
    main()


