#!/usr/bin/env python
# TODO: check first if argument parameters are of the right type if given.
import os
import sys
import argparse
import logging
import rpy2.robjects as robjects
import utils

rscript = ''
R = robjects.r


def run_rscript(command=None):
    """Run R command, log it, append to rscript"""
    global rscript
    if not command:
        return
    logging.debug(command)
    rscript += '{}\n'.format(command)
    R(command)


def plot_transcript(rdata_load='Metagene.rda', transcript_name='',
                    transcript_length='27', transcript_cap='',
                    html_file='Plot-ribosome-profile.html',
                    output_path=os.getcwd()):
    """Plot ribosome profile for a given transcript. """
    options = {}
    for key, value, rtype, rmode in (
            ('transcript_name', transcript_name, 'str', None),
            ('transcript_length', transcript_length, 'int', 'charvector'),
            ('transcript_cap', transcript_cap, 'int', None)):
        options[key] = utils.process_args(value, ret_type=rtype, ret_mode=rmode)

    run_rscript('library(riboSeqR)')
    run_rscript('load("{}")'.format(rdata_load))

    html = '<h2>Plot ribosome profile - results</h2><hr>'
    if len(transcript_name):
        cmd_args = (
            '"{transcript_name}", main="{transcript_name}",'
            'coordinates=ffCs@CDS, riboData=riboDat,'
            'length={transcript_length}'.format(**options))
        if transcript_cap:
            cmd_args += ', cap={transcript_cap}'.format(**options)
        cmd = 'plotTranscript({})'.format(cmd_args)

        for fmat in ('pdf', 'png'):
            run_rscript('{0}(file="{1}")'.format(
                fmat, os.path.join(
                    output_path, 'Ribosome-profile-plot.{}'.format(fmat))))
            run_rscript(cmd)
            run_rscript('dev.off()')

        html += ('<p>Selected ribosome footprint length: <strong>{0}</strong>'
                 '</p><p><img border="1" src="Ribosome-profile-plot.png" '
                 'alt="Ribosome profile plot" /><br>'
                 '<a href="Ribosome-profile-plot.pdf">PDF version</a>'
                 '</p>'.format(transcript_length))
    else:
        msg = 'No transcript name was provided. Did not generate plot.'
        html += '<p>{}</p>'.format(msg)
        logging.debug(msg)

    logging.info('\n{:#^80}\n{}\n{:#^80}\n'.format(
        ' R script for this session ', rscript, ' End R script '))

    with open(os.path.join(output_path, 'ribosome-profile.R'), 'w') as r:
        r.write(rscript)

    html += ('<h4>R script for this session</h4>'
             '<p><a href="ribosome-profile.R">ribosome-profile.R</a></p>')

    with open(html_file, 'w') as f:
        f.write(html)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot Ribosome profile')

    # required arguments
    flags = parser.add_argument_group('required arguments')
    flags.add_argument('--rdata_load', required=True,
                       help='Saved riboSeqR data from Step 2')
    flags.add_argument('--transcript_name', required=True,
                       help='Name of the transcript to be plotted')
    flags.add_argument(
        '--transcript_length', required=True,
        help='Size class of ribosome footprint data to be plotted',
        default='27')
    flags.add_argument(
        '--transcript_cap', required=True,
        help=('Cap on the largest value that will be plotted as an abundance '
              'of the ribosome footprint data'))
    parser.add_argument('--html_file', help='HTML file with reports')
    parser.add_argument('--output_path', help='Directory to save output files')
    parser.add_argument('--debug', help='Produce debug output',
                        action='store_true')

    args = parser.parse_args()
    if args.debug:
        level = logging.DEBUG
    else:
        level = logging.INFO

    logging.basicConfig(format='%(levelname)s - %(message)s', level=level,
                        stream=sys.stdout)
    logging.debug('Supplied Arguments\n{}\n'.format(vars(args)))

    if not os.path.exists(args.output_path):
        os.mkdir(args.output_path)

    plot_transcript(rdata_load=args.rdata_load,
                    transcript_name=args.transcript_name,
                    transcript_length=args.transcript_length,
                    transcript_cap=args.transcript_cap,
                    html_file=args.html_file, output_path=args.output_path)
    logging.debug('Done!')