import sys
import os
import glob
import re
import click
import numpy as np
import pysam
import random
import string
from collections import defaultdict
import logging
from typing import *
import velocyto as vcy
from ._run import _run

# logging.basicConfig(stream=sys.stdout, format='%(asctime)s - %(levelname)s - %(message)s', level=logging.DEBUG)


def id_generator(size: int=6, chars: str=string.ascii_uppercase + string.digits) -> str:
    return ''.join(random.choice(chars) for _ in range(size))


@click.command(short_help="Runs the velocity analysis outputing a loom file")
@click.argument("bamfilepath", nargs=1, required=True,
                type=click.Path(exists=True,
                                file_okay=True,
                                dir_okay=False,
                                readable=True,
                                resolve_path=True))
@click.argument("dropest-out", nargs=1, required=True,
                type=click.Path(exists=True,
                                file_okay=True,
                                dir_okay=False,
                                readable=True,
                                resolve_path=True))
@click.option("--corrected-output", "-o",
              help="""(Optional) The file output of the output bam file. Otherwise the file will be outputed in the same folder of the input with the prefix `correct_`""",
              default=None,
              type=click.Path(exists=False,
                              resolve_path=True,
                              file_okay=True,
                              dir_okay=False))
def dropest_bc_correct(bamfilepath: str, dropest_out: str, corrected_output: str) -> None:
    """Using the output of DropEst:
    (1) Corrects barcodes directly in the bam file
    (2) Produces a valid barcodes list

    BAMFILEPATH - bam file with sorted reads obtained running DropEst

    DROPEST-OUT - R dump `rds` file generated by DropEst
    """

    try:
        import rpy2.robjects as ro
        from velocyto.r_interface import convert_r_obj
    except:
        ImportError("A problem was encountered importing rpy2. To run this `velocyto tools` rpy2 and R need to be installed (the use conda is reccomended)")

    # bamfilepath = sys.argv[1]
    # filename = os.path.join(parentpath, f"{bamfilename.split('_')[0]}_dropEst.rds")
    parentpath, bamfilename = os.path.split(bamfilepath)
    filename = dropest_out
    logging.info(f"Loading `merge_targets` from {filename} using R / rpy2")
    mapping = convert_r_obj(ro.r(f"rds <- readRDS('{filename}'); rds$merge_targets"))  # a dict

    output_path = os.path.join(parentpath, f"barcodes_{bamfilename.split('_')[0]}.tsv")
    logging.info(f"Generating {output_path}")
    with open(output_path, "w") as fout:
        unique_bcs = set(list(mapping.values()))
        str_ = '\n'.join(list(unique_bcs))
        fout.write(str_)

    infile = pysam.AlignmentFile(bamfilepath, mode="rb")
    if corrected_output is None:
        bam_out_path = os.path.join(parentpath, f"correct_{bamfilename}")
    else:
        bam_out_path = corrected_output

    outfile = pysam.AlignmentFile(bam_out_path, mode="wb", template=infile)

    for read in infile:
        # read: pysam.AlignedRead
        cb = read.get_tag("CB")
        if cb in mapping:
            read.set_tag("CB", mapping[cb], value_type="Z")
        outfile.write(read)
    infile.close()
    outfile.close()
    logging.info("Done")

    return
