# -*- coding: utf-8 -*-
"""Script for generating exon expression counts."""

# pylint: disable=wildcard-import,redefined-builtin,unused-wildcard-import
from __future__ import absolute_import, division, print_function
from builtins import *
# pylint: enable=wildcard-import,redefined-builtin,unused-wildcard-import

import argparse
import logging

from pathlib2 import Path

import imfusion
from imfusion.build.indexers import Reference
from imfusion.expression.counts import generate_exon_counts

FORMAT = "[%(asctime)-15s] %(message)s"
logging.basicConfig(format=FORMAT, level=logging.INFO)

# Map for translating strand options to numeric values used by featureCounts.
STRANDED_MAP = {'unstranded': 0, 'stranded': 1, 'reverse': 2}


def main():
    """Main function for imfusion-expression."""

    args = parse_args()

    # Assemble featureCounts arguments from options.
    feature_counts_kws = {'-s': STRANDED_MAP[args.stranded]}

    if args.paired:
        feature_counts_kws['-p'] = True

    # Generate exon counts.
    bam_path = args.sample_dir / 'alignment.bam'
    reference = Reference(args.reference)

    exon_counts = generate_exon_counts(
        [bam_path],
        gtf_path=reference.exon_gtf_path,
        names={str(bam_path): bam_path.parent.name},
        extra_kws=feature_counts_kws)

    # Write output.
    output = args.output or args.sample_dir / 'expression.txt'
    exon_counts.to_csv(str(output), sep='\t', index=True)


def parse_args():
    """Parses arguments for imfusion-expression."""

    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--version',
        action='version',
        version='IM-Fusion ' + imfusion.__version__)

    parser.add_argument(
        '--sample_dir',
        type=Path,
        required=True,
        help=('Path to sample directory (typically the output '
              'of imfusion-insertions).'))

    parser.add_argument(
        '--reference',
        type=Path,
        required=True,
        help='Path to the index of the augmented reference '
        'generated by im-fusion build.')

    parser.add_argument('--output', type=Path, default=None)

    parser.add_argument(
        '--paired',
        default=False,
        action='store_true',
        help='Generate counts by counting fragments '
        'instead of reads (for paired-end data).')

    parser.add_argument(
        '--stranded',
        default='unstranded',
        choices={'unstranded', 'stranded', 'reverse'},
        help='Strandedness of the RNA-seq data.')

    return parser.parse_args()
