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

import argparse
import pysam
import sys
from datetime import datetime
import singlecellmultiomics
from singlecellmultiomics.fragment import Fragment
from singlecellmultiomics.bamProcessing.bamFunctions import get_read_group_from_read, sorted_bam_file, write_program_tag

def set_read_group( in_bam_path, out_bam_path, id:str, pl:str, lb:str, sm:str, pu:str, threads=4, sm2reads=False  ):
    """
    Set read group format of bam file

    Args:
        in_bam_path(str) : bam file to read

        in_bam_path(str) : bam file to write

        format(int) : formatting mode, see `singlecellmultiomics.fragment.Fragment.get_read_group`

        threads(int) : Amount of decompression threads for reading

    """

    """
    read_groups(set/dict) : set or dictionary which contains read groups. The dictionary should have the format { read_group_id (str)
            { 'ID': ID, 'LB':library,
            'PL':platform,
            'SM':sampleLib,
            'PU':readGroup }
    """

    read_groups =     {id:{ 'ID': id, 'LB':lb,
        'PL':pl,
        'SM':sm,
        'PU':pu }}

    with pysam.AlignmentFile(in_bam_path, threads = threads) as input_bam:

        input_header = input_bam.header.as_dict()

        # Write provenance information to BAM header
        write_program_tag(
            input_header,
            program_name='bamReadGroupFormat',
            command_line=" ".join(
                sys.argv),
            version=singlecellmultiomics.__version__,
            description=f'SingleCellMultiOmics read group formatting, executed at {datetime.now().strftime("%d/%m/%Y %H:%M:%S")}')

        with sorted_bam_file(out_bam_path, header=input_header, read_groups=read_groups, input_is_sorted=True) as out:
            print('Started writing')
            for read in input_bam:
                rg_id = id
                read.set_tag('RG',rg_id)
                if sm2reads:
                    read.set_tag('SM',sm)
                out.write(read)


if __name__=='__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='Set read group sample platform unit etc.')
    argparser.add_argument('bamin', type=str)


    argparser.add_argument(
        '-id',
        type=str,
        required = True,
        help="Read group id")

    argparser.add_argument(
        '-pl',
        type=str,
        default='ILLUMINA',
        help="Read group platform")

    argparser.add_argument(
        '-pu',
        type=str,
        default='LANE',
        help="Read group platform unit")

    argparser.add_argument(
        '-sm',
        type=str,
        required=True,
        help="Read group sample name")

    argparser.add_argument(
        '-lb',
        type=str,
        default='mix',
        help="Library")


    argparser.add_argument(
        '-t',
        type=int,
        default=4,
        help="Threads")


    argparser.add_argument(
        '--sm2reads',
        action='store_true',
        help="Also write the supplied SM tag to all reads")

    argparser.add_argument('-o', type=str, help="output bam file", required=True)
    args = argparser.parse_args()
    set_read_group(
        args.bamin,
        args.o,
        id = args.id,
        pl = args.pl,
        lb = args.lb,
        sm = args.sm,
        pu = args.pu,
        sm2reads= args.sm2reads,
        threads = args.t)