#!/usr/bin/python

## Global modules
import os
from os.path import join
import logging
import shutil
import pysam

## Local modules
import mglobals
import helpers
import snp2gene

log = logging.getLogger('pipeline')


@helpers.log_func
def my_tophat():

    if mglobals.original:
        log.info('Aligning reads to the original reference fasta')
        fastas = mglobals.original_fastas
    else:
        log.info('Aligning reads to alternate reference fastas')
        fastas = mglobals.alternate_fastas

    @helpers.multiprocess(zip(mglobals.samples_list, fastas))
    def tophat_call(sample, ref_fasta):

        if mglobals.original:
            os.chdir(join(mglobals.original_path, sample))
        else:
            os.chdir(join(mglobals.alternate_path, sample))

        ref_fasta_base = ref_fasta.split('.')[0]

        mismatches = '2'
        number_of_samples = len(mglobals.samples_list)
        threads_per_sample = mglobals.cpu_count//number_of_samples

        threads = str(threads_per_sample)
        log.info('threads per sample ' + threads)

        log.info('tophat: aligning sample {} with ref fasta {}'.format(sample, ref_fasta))
        tophat_params = ['nice', '-n', '5',
                         'tophat',
                         '-p', threads,
                         '-G', mglobals.dros_gtf,
                         '--transcriptome-index=../transcriptome_data/known',
                         '-N', mismatches,
                         '--b2-L', '20',
                         '--b2-N', '1',
                         '--read-edit-dist', mismatches,
                         '-o', (sample + '_thout'),
                         '--no-novel-juncs',
                         ref_fasta_base,
                         join(mglobals.samples_path, (sample + '.fastq'))]
        helpers.sub_call(tophat_params)

        log.info('tophat: finished analyzing sample: {} with ref fasta: {}'.format(sample, ref_fasta))

    # Copy transcriptome index to original_path and alternate_path
    for path in [mglobals.original_path, mglobals.alternate_path]:
        if not os.path.exists(join(path, 'transcriptome_data')):
            log.info('Linking transcriptome data to ' + path)
            os.symlink(mglobals.dros_gtf_index, join(path, 'transcriptome_data'))

    tophat_call()


@helpers.log_func
def my_alignment_filter():

    @helpers.multiprocess(mglobals.samples_list)
    def filter_call(sample):

        if mglobals.original:
            os.chdir(join(mglobals.original_path, sample, (sample + '_thout')))
        else:
            os.chdir(join(mglobals.alternate_path, sample, (sample + '_thout')))

        log.info('Filtering aligned reads for: ' + sample)

        # Index each bamfile
        if not os.path.exists('accepted_hits.bam.bai'):
            pysam.index('accepted_hits.bam')

        # Sort by NH flag
        raw_reads = pysam.Samfile('accepted_hits.bam', 'rb')
        filter_reads = pysam.Samfile('filter.bam', 'wb', template=raw_reads)
        for read in raw_reads.fetch():
            if ('NH', 1) in read.tags:
                filter_reads.write(read)
        raw_reads.close()
        filter_reads.close()

        pysam.index('filter.bam')

    filter_call()


@helpers.log_func
def my_pileup(out_file_extension='.mpileup'):

    if mglobals.original:
        fastas = mglobals.original_fastas
    else:
        fastas = mglobals.alternate_fastas

    @helpers.multiprocess(zip(mglobals.samples_list, fastas))
    def pileup_call(sample, ref_fasta):