#!python
import logging
logging.getLogger().setLevel(logging.INFO)  # override config.py's DEBUG before any scTE imports
import pandas as pd
import multiprocessing
from functools import partial
import os, sys, glob, datetime, time, gzip, shutil
from importlib.metadata import version
import argparse
import collections
from math import log
sys.path.append(os.path.join(os.path.split(sys.argv[0])[0], '../'))
from scTE.miniglbase import genelist, glload, location
from scTE.annotation import anno_gtf
from scTE.base import *
from scTE.base import set_compress_mode, compress_ext, compress_cmd, decompress_cmd, _p, xopen
import scTE.base as _scTE_base

__version__ = version("scte-quant")

def prepare_parser():
    desc = "Quantifying transposable element (TEs) expression from (single-cell) RNA sequencing data"

    exmp = "Example: scTE <-i scRNA.sorted.bam> <-o out> [--min_genes 200] [--min_counts 400] [-p 4] <-x mm10.exclusive.idx>"

    parser = argparse.ArgumentParser(prog='scTE',description=desc, epilog=exmp)

    optional = parser._action_groups.pop()

    optional.add_argument('--min_genes', dest='genenumber',metavar='INT', type=int,default=200,
                        help='Minimum number of genes expressed required for a cell to pass filtering. Default: 200')

    optional.add_argument('--min_counts', dest='countnumber',metavar='INT', type=int,
                        help='Minimum number of counts required for a cell to pass filtering. Default: 2*min_genes')

    optional.add_argument('--expect-cells', dest='cellnumber',metavar='INT', type=int,  default=10000,
                        help='Expected number of cells. Default: 10000')

    optional.add_argument('-f','--format', metavar='input file format', dest='format', type=str, nargs='?', default='BAM', choices=['BAM','SAM'],
                        help='Input file format: BAM or SAM. DEFAULT: BAM')

    optional.add_argument('-CB', dest='CB', type=str, nargs='?', default='CR', choices=['CR','CB','False'],
                        help='Set to false to ignore for cell barcodes, it is useful for SMART-seq. If you set CB=False, it also will set UMI=False by default, Default: CR')

    optional.add_argument('-UMI', dest='UMI', type=str, nargs='?', default='UR', choices=['UR','UB','False'],
                        help='Set to false to ignore for UMI, it is useful for SMART-seq. Default: UR')

    optional.add_argument('--keeptmp', dest='keeptmp', action='store_true', default=False,
                        help='Keep the _scTEtmp file, which is useful for debugging. Default: off')

    optional.add_argument('--hdf5', dest='hdf5', action='store_true', default=False,
                        help='Save the output as .h5ad formatted file instead of csv file. Default: off')

    optional.add_argument('-p','--thread', metavar='INT', dest='thread', type=int, default=1,
                        help='Number of threads to use, Default: 1')

    optional.add_argument('--max-chunk-lines', metavar='INT', dest='max_chunk_lines', type=int, default=0,
                        help='Maximum lines per chunk for large chromosome BEDs in align phase. '
                             '0 = adaptive (targets 2x thread count tasks). Default: 0')
    
    optional.add_argument('--intermediate-compression', dest='compress', type=str, nargs='?', default='zstd',
                         choices=['gzip', 'zstd', 'none'],
                         help='Compression for intermediate BED files. Default: zstd')

    optional.add_argument('--use-polars-bio', dest='use_polars_bio', type=str, default='auto',
                         choices=['auto', 'yes', 'no'],
                         help='Use polars-bio for accelerated BAM processing. auto = use if installed. Default: auto')

    optional.add_argument('--verbose', dest='verbose', action='store_true', default=False,
                        help='Show detailed progress (per-chromosome info). Default: off')
    optional.add_argument('-q','--quiet', dest='quiet', action='store_true', default=False,
                        help='Suppress non-critical output. Default: off')

    optional.add_argument('-v','--version', action='version', version=f'%(prog)s {__version__}')

    required = parser.add_argument_group('required arguments')

    required.add_argument('-i','--input', dest='input', type=str, nargs='+', required=True,
                        help='Input file: BAM/SAM file from CellRanger or STARsolo, the file must be sorted by chromosome position')

    required.add_argument('-x', dest='annoglb',nargs='+', required=True,
                        help='The filename of the index for the reference genome annotation.')

    required.add_argument('-o','--out', dest='out', nargs='?', required=True, help='Output file prefix')

    parser._action_groups.append(optional)
    optional = parser.add_argument_group('optional arguments')
    optional

    return parser

def main():
    """Start scTEs......parse options......"""

    timestart=datetime.datetime.now()
    args=read_opts(prepare_parser())
    set_compress_mode(args.compress)

    # --- Verbosity control ---
    # config.py sets root logger to DEBUG at import time; override it here.
    if args.quiet:
        logging.getLogger().setLevel(logging.WARNING)
        logging.getLogger('glbase3').setLevel(logging.WARNING)
    elif args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
        logging.getLogger('glbase3').setLevel(logging.INFO)
        # Verbose mode: restore polars-bio loggers (suppressed by default in base.py)
        logging.getLogger("polars_bio").setLevel(logging.INFO)
        logging.getLogger("polars_bio.context").setLevel(logging.INFO)
        logging.getLogger("polars_bio.io").setLevel(logging.INFO)
        logging.getLogger("datafusion_bio_function_ranges").setLevel(logging.INFO)
        logging.getLogger("datafusion_bio_function_ranges.session_context").setLevel(logging.INFO)
    else:
        logging.getLogger().setLevel(logging.INFO)
        logging.getLogger('glbase3').setLevel(logging.WARNING)  # suppress per-thread "Loaded" messages
    logging.getLogger('h5py').setLevel(logging.WARNING)  # suppress codec registration spam
    # ---

    log_fh = logging.FileHandler('%s.log' % args.out, mode='w')
    log_fh.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S'))
    logging.getLogger().addHandler(log_fh)
    # ---

    info = args.info
    error = args.error

    assert sys.version_info >= (3, 6), 'Python >=3.6 is required'

    info(args.argtxt + "\n")

    outname = args.out.split('/')[-1:][0]

    info("Loading the genome annotation index... %s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
    allelement, chr_list, all_annot, glannot = read_anno(filename=outname, annoglb=args.annoglb[0]) #genome=args.genome
    if args.verbose:
        args.debug(sorted(chr_list))
    info("Finished loading the genome annotation index... %s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

    info("Processing BAM/SAM files ...%s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

    if len(args.input) == 1 and ',' in args.input[0]:
        args.input=args.input[0].split(',')

    # --- Activate polars-bio path if requested ---
    _use_polars_bio = (args.use_polars_bio == 'yes' or 
                       (args.use_polars_bio == 'auto' and _scTE_base._HAS_POLARS_BIO))
    if _use_polars_bio:
        _scTE_base.set_use_polars_bio(True)
        info("Using polars-bio accelerated path")

    if _use_polars_bio:
        # ------ polars-bio fast path ------
        os.makedirs('%s_scTEtmp/o1' % outname, exist_ok=True)
        for k in args.input:
            check_cb_umi(filename=k, out=outname, CB=args.CB, UMI=args.UMI)
        info("Input BAM file appears to be valid")
        info("Done BAM/SAM files validation ...%s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

        info("Building barcode whitelist ...%s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
        align_t0 = time.time()
        if len(args.input) > 1:
            # Multi-sample: process each BAM separately
            all_res = {}
            all_stats = []
            for bam_path in args.input:
                sid = os.path.basename(bam_path).replace('.bam', '')
                whitelist = build_whitelist_polars(
                    bam_path, args.CB, args.genenumber,
                    args.countnumber, sample_name=sid)
                res, stats = align_polars(
                    bam_path, _scTE_base._SHARED_ANNOT_DF, whitelist,
                    args.CB, args.UMI, sample_name=sid)
                # Merge multi-sample results (barcodes already have sample prefix)
                for bc, genes in res.items():
                    if bc not in all_res:
                        all_res[bc] = defaultdict(int)
                    for gene, cnt in genes.items():
                        all_res[bc][gene] += cnt
                all_stats.extend(stats)
            res = all_res
            align_stats = all_stats
        else:
            bam_path = args.input[0]
            whitelist = build_whitelist_polars(
                bam_path, args.CB, args.genenumber,
                args.countnumber, sample_name=outname)
            info("Whitelist: %d barcodes" % len(whitelist))
            info("Fetching from the annotation index ...%s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
            res, align_stats = align_polars(
                bam_path, _scTE_base._SHARED_ANNOT_DF, whitelist,
                args.CB, args.UMI, sample_name=outname)
        align_wall = time.time() - align_t0
        info("Done fetching ...%s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

        # Write o3 BED file for count_expression() and compress to o4
        os.makedirs('%s_scTEtmp/o3' % outname, exist_ok=True)
        o3_file = '%s_scTEtmp/o3/%s.all.bed' % (outname, outname)
        with open(o3_file, 'w') as oh:
            for bc in sorted(res):
                for gene in sorted(res[bc]):
                    oh.write('%s\t%s\t%s\n' % (bc, gene, res[bc][gene]))
        os.makedirs('%s_scTEtmp/o4' % outname, exist_ok=True)
        o4 = _p('%s_scTEtmp/o4/%s.bed.gz' % (outname, outname))
        with open(o3_file, 'rb') as src:
            with open(o4, 'wb') as dst:
                subprocess.run(compress_cmd().split(), stdin=src, stdout=dst, check=True)
        os.unlink(o3_file)
        info("Done fetching... %s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

    else:
        # ------ Original (non-polars-bio) path ------
        os.makedirs('%s_scTEtmp/o1' % outname, exist_ok=True)

        for k in args.input:
            check_cb_umi(filename=k,out=outname,CB=args.CB,UMI=args.UMI)
        info("Input SAM/BAM file appears to be valid")

        if len(args.input) > 1:
            info('Using parabam2bed as more than 1 input BAM')
            n_files = len(args.input)
            pool_size = min(args.thread, n_files)
            per_worker = max(1, args.thread // n_files)
            with multiprocessing.Pool(processes=pool_size) as pool:
                partial_work = partial(para_bam2bed, CB=args.CB, UMI=args.UMI, out=outname, num_threads=per_worker)
                pool.map(partial_work, args.input)
            out1 = _p('%s_scTEtmp/o1/%s.bed.gz' % (outname, outname))
            subprocess.run('%s %s_scTEtmp/o0/*.bed* | %s > %s' % (decompress_cmd(), outname, compress_cmd(), out1),
                            shell=True, check=True)
        
        else:
            args.debug('%s %s\n' % (args.CB, args.UMI))
            bam2bed(args.input[0], args.CB, args.UMI, outname, args.thread)
        info("Done BAM/SAM files processing ...%s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

        info("Splitting ...%s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
        if args.thread == 1: #Single thread path, mainly
            info('Executing single thread path')
            whitelist = split_all_chrs(chr_list, filename=outname, genenumber=args.genenumber, countnumber=args.countnumber, UMI=args.UMI)
        else:
            info('Executing multiple thread path with %s threads' % args.thread)
            with multiprocessing.Pool(processes=args.thread) as pool:
                partial_work = partial(split_chr, filename=outname, CB=args.CB, UMI=args.UMI)
                pool.map(partial_work, chr_list)
            whitelist = filter_crs(filename=outname, genenumber=args.genenumber, countnumber=args.countnumber)

        info("Finished processing sample files %s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

        info("Fetching from the annotation index... %s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
        align_t0 = time.time()
        align_stats = []

        if args.thread == 1: #Single thread path
            for chrom in chr_list:
                args.debug("align: processing chromosome %s" % chrom)
                result = align(chr_spec=chrom, filename=outname, all_annot=None, glannot=glannot, whitelist=whitelist)
                align_stats.append(result)

        else: # Multiprocessing path:
            # --- Cache the index before forking workers ---
            cache_index(glannot, all_annot)

            # --- Split large chromosome BEDs into chunks for load balancing ---
            align_tasks = split_bed_for_align(outname, chr_list, max_lines=args.max_chunk_lines, n_threads=args.thread)

            n_chunks = sum(1 for t in align_tasks if '_chunk' in t)
            if n_chunks:
                info("Align: %d chromosome tasks + %d chunk tasks = %d total (pool size=%d)" %
                     (len(chr_list), n_chunks, len(align_tasks), args.thread))

            with multiprocessing.Pool(processes=args.thread) as pool:
                partial_work = partial(align, filename=outname, all_annot=all_annot, glannot=None, whitelist=whitelist)
                for result in pool.imap_unordered(partial_work, align_tasks):
                    align_stats.append(result)

        align_wall = time.time() - align_t0

        os.makedirs('%s_scTEtmp/o4' % outname, exist_ok=True)
        out4 = _p('%s_scTEtmp/o4/%s.bed.gz' % (outname, outname))
        o3_dir = '%s_scTEtmp/o3' % outname
        o3_pat = os.path.join(o3_dir, '%s.*.bed*' % outname)
        tmp4 = _p('%s_scTEtmp/o4/%s.bed.raw' % (outname, outname))
        with open(tmp4, 'w') as tmp:
            for f in sorted(glob.glob(o3_pat)):
                with xopen(f, 'rt') as src:
                    for line in src:
                        tmp.write(line)
        with open(tmp4, 'rb') as src:
            with open(out4, 'wb') as dst:
                subprocess.run(compress_cmd().split(), stdin=src, stdout=dst, check=True)
        os.unlink(tmp4)
        info("Done fetching... %s \n"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

    # --- Common post-processing ---

    # Write per-chromosome alignment statistics to .stat file
    read_counts = get_chr_read_counts(outname, chr_list)
    write_stat_file(outname, chr_list, read_counts, align_stats, align_wall)

    info("Calculating expression... %s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
    len_res, genenumber, filename = count_expression(filename=outname, allelement=allelement, genenumber=args.genenumber, cellnumber=args.cellnumber, hdf5=args.hdf5)
    if args.hdf5:
        info('Detect {0} cells expressed at least {1} genes, results output to {2}.h5ad'.format(len_res, genenumber, filename))
    else:
        info('Detect {0} cells expressed at least {1} genes, results output to {2}.csv'.format(len_res, genenumber, filename))
    
    info("Finished calculating expression %s"%(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))

    if args.keeptmp:
        pass
    else:
        shutil.rmtree('%s_scTEtmp' % outname, ignore_errors=True)

    timeend = datetime.datetime.now()
    info("Done with %s\n" % timediff(timestart,timeend))

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.stderr.write("User interrupt !\n")
        sys.exit(0)


