#!/usr/bin/env python

# -*- coding: utf-8 -*-

import pysam

import argparse

def convert_scmo_to_cellranger(scmo_bam, cellranger_bam, n_threads=4):

    """

    Convert a SCMO bam file to a cellranger bam file

    Adds these tags:

        BC to CB

        RX to UB

    """

    with pysam.AlignmentFile(scmo_bam,threads=n_threads) as al:

        with pysam.AlignmentFile(cellranger_bam,  'wb', header=al.header, threads=n_threads) as ao:

            for read in al:

                if read.has_tag('BC'):

                    read.set_tag('CB',read.get_tag('BC'))

                if read.has_tag('RX'):

                    read.set_tag('UB',read.get_tag('RX'))

                ao.write(read)

    pysam.index(cellranger_bam)

def convert_scmo_to_dropseq(scmo_bam, dropseq_bam, n_threads=4):

    """

    Convert a SCMO bam file to a dropseq bam file

    Adds these tags:

        BC to XC

        RX to XM

    """

    with pysam.AlignmentFile(scmo_bam,threads=n_threads) as al:

        with pysam.AlignmentFile(dropseq_bam,  'wb', header=al.header, threads=n_threads) as ao:

            for read in al:

                if read.has_tag('BC'):

                    read.set_tag('XC',read.get_tag('BC'))

                if read.has_tag('RX'):

                    read.set_tag('XM',read.get_tag('RX'))

                ao.write(read)

    pysam.index(dropseq_bam)

if __name__ == '__main__':

    argparser = argparse.ArgumentParser(

        formatter_class=argparse.ArgumentDefaultsHelpFormatter,

        description="""Convert SCMO bam file to cellranger or dropseq bam file

    """)

    argparser.add_argument('inputbamfile', type=str)

    argparser.add_argument('convertedbamfile', type=str)

    og = argparser.add_argument_group("Output")

    #og.add_argument('-bed', type=str, help='Bed file to write methylation calls to')

    og.add_argument('-fmt', type=str, help='Format to convert to cellranger/dropseq', required=True)

    args = argparser.parse_args()

    if args.fmt == 'cellranger':

        convert_scmo_to_cellranger(args.inputbamfile, args.convertedbamfile)

    elif args.fmt == 'dropseq':

        convert_scmo_to_dropseq(args.inputbamfile, args.convertedbamfile)