
"""
Utility functions for regenerating our test data starting from h5 files.
"""

import multiprocessing
import subprocess
from collections import defaultdict
from functools import partial as P
import tempfile
import logging
import shutil
import random
import time
import re
import os.path as op
import os
import sys

import pysam

from pbcore.io import openDataFile


def import_movie_metadata(metadata_file, ds_file):
    cmd = [
        "movie-metadata-to-dataset",
        metadata_file,
        ds_file
    ]
    logging.info("Converting %s" % metadata_file)
    logfile = op.splitext(ds_file)[0] + "_import.log"
    with open(logfile, "w") as f:
        assert subprocess.call(cmd, stdout=f) == 0
    assert op.isfile(ds_file)
    return ds_file


def __import_h5_reads(extension, read_type, file_name, extra_args=()):
    """
    Convert a single h5 file to <read_type>.bam, keeping the movie name and
    file number intact.
    """
    base = re.sub(extension, "", file_name)
    bam_file = base + ".%s.bam" % read_type
    args1 = ["bax2bam", "-o", base, file_name] + list(extra_args)
    logging.info(" ".join(args1))
    logfile = base + "_bax2bam.log"
    with open(logfile, "w") as f:
        assert subprocess.call(args1, stdout=f) == 0
    assert op.isfile(bam_file) and op.isfile(bam_file + ".pbi")
    return bam_file


import_subreads = P(__import_h5_reads, "\.bax\.h5", "subreads")
import_ccs = P(__import_h5_reads, "\.ccs\.h5", "ccs", extra_args=["--ccs"])


def index_bam(file_name):
    """Run 'samtools index' and 'pbindex' on a BAM file."""
    def _remove_index(ext):
        idx_file = file_name + "." + ext
        if op.exists(idx_file):
            os.remove(idx_file)
    base = op.splitext(file_name)[0]
    _remove_index("bai")
    args2 = ["samtools", "index", file_name]
    logging.info(" ".join(args2))
    logfile = base + "_samtools.log"
    with open(logfile, "w") as f:
        assert subprocess.call(args2, stdout=f) == 0
    _remove_index("pbi")
    args3 = ["pbindex", file_name]
    logging.info(" ".join(args3))
    logfile = base + "_pbindex.log"
    with open(logfile, "w") as f:
        assert subprocess.call(args3, stdout=f) == 0
    return file_name


def align_subreads(bam_file, ref_file, out_file):
    """
    Run pbalign to generate an aligned BAM file.
    """
    prefix = re.sub(".subreads.bam", "", bam_file)
    tmp_out_file = tempfile.NamedTemporaryFile(suffix=".bam").name
    args = ["pbalign", bam_file, ref_file, tmp_out_file]
    logfile = prefix + "_pbalign.log"
    with open(logfile, "w") as f:
        assert subprocess.call(args, stdout=f) == 0
    if op.exists(out_file):
        os.remove(out_file)
    if op.exists(out_file + ".pbi"):
        os.remove(out_file + ".pbi")
    if op.exists(out_file + ".bai"):
        os.remove(out_file + ".bai")
    shutil.move(tmp_out_file, out_file)
    shutil.move(tmp_out_file + ".pbi", out_file + ".pbi")
    shutil.move(tmp_out_file + ".bai", out_file + ".bai")
    return out_file


def _func_wrapper(func, args):
    try:
        return func(*args)
    except Exception as e:
        return e

_align_subreads_wrapper = P(_func_wrapper, align_subreads)
_import_subreads_wrapper = P(_func_wrapper, import_subreads)
_import_ccs_wrapper = P(_func_wrapper, import_ccs)
_index_bam_wrapper = P(_func_wrapper, index_bam)
_import_movie_metadata_wrapper = P(_func_wrapper, import_movie_metadata)


def _run_functions_parallel(func, args, nproc):
    if len(args) == 0:
        return []
    if isinstance(args[0], basestring):
        args = [(arg,) for arg in args]
    t1 = time.time()
    p = multiprocessing.Pool(processes=max(nproc, 1))
    results = p.map(func, args)
    p.close()
    t2 = time.time()
    logging.info("Finished %d tasks in %.1fs" % (len(args), t2 - t1))
    for args_, result in zip(args, results):
        if isinstance(result, Exception):
            logging.error(result)
            logging.error("Failed processing %s" % str(args_))
            raise result
    return results

align_subreads_parallel = P(_run_functions_parallel, _align_subreads_wrapper)
index_bam_parallel = P(_run_functions_parallel, _index_bam_wrapper)
import_subreads_parallel = P(_run_functions_parallel, _import_subreads_wrapper)
import_ccs_parallel = P(_run_functions_parallel, _import_ccs_wrapper)
import_movie_metadata_parallel = P(_run_functions_parallel,
    _import_movie_metadata_wrapper)


def extract_small_dataset_sample(file_name, n_reads=0, n_zmws=0,
                                 output_file=None,
                                 randomize=False):  # FIXME
    """
    From the input dataset, extract the first N reads or ZMWs from each .bam
    file and write to an identically named .bam in the current directory.
    Used to generate micro-datasets for very fast testing of pbsmrtpipe.
    """
    assert ([n_reads, n_zmws].count(0) == 1)
    if n_reads == 0:
        n_reads = sys.maxint
    dataset_type = None
    xml_files = []
    if output_file is None:
        output_file = op.basename(file_name)
    if op.abspath(output_file) == op.abspath(file_name):
        output_file = op.splitext(output_file)[0] + "_tiny.xml"
    with openDataFile(file_name) as ds:
        dataset_type = type(ds)
        for bam in ds.resourceReaders():
            logging.info("processing %s" % op.basename(bam.filename))
            bam_output = op.splitext(output_file)[0] + ".bam"
            with pysam.AlignmentFile(bam_output, "wb",
                                     template=bam.peer) as out:
                zmws = set([])
                n_file_reads = n_file_zmws = 0
                if randomize:  # FIXME this should probably be the default
                    if n_zmws > 0:
                        zmw_dict = defaultdict(list)
                        for i_read, zmw in enumerate(bam.holeNumber):
                            zmw_dict[zmw].append(i_read)
                        have_zmws = set()
                        have_reads = set()
                        zmws = zmw_dict.keys()
                        while True:
                            i_zmw = random.randint(0, len(zmws) - 1)
                            if not zmws[i_zmw] in have_zmws:
                                for i_read in zmw_dict[zmws[i_zmw]]:
                                    assert not i_read in have_reads
                                    out.write(bam[i_read].peer)
                                    have_reads.add(i_read)
                                    n_file_reads += 1
                                have_zmws.add(zmws[i_zmw])
                                n_file_zmws += 1
                            if n_file_zmws == n_zmws:
                                break
                    else:
                        have_reads = set()
                        while True:
                            i_read = random.randint(0, len(bam) - 1)
                            if not i_read in have_reads:
                                out.write(bam[i_read].peer)
                                have_reads.add(i_read)
                                n_file_reads += 1
                            if n_file_reads == n_reads:
                                break
                else:
                    for read in bam:
                        if n_zmws:
                            zmws.add(read.HoleNumber)
                        else:
                            n_file_reads += 1
                        if len(zmws) > n_zmws or n_file_reads > n_reads:
                            break
                        else:
                            out.write(read.peer)
                            n_file_reads += 1
                logging.info("Wrote %d reads to file %s" % (n_file_reads,
                                                            bam_output))
            subprocess.call(["samtools", "index", bam_output])
            subprocess.call(["pbindex", bam_output])
            ds_out = op.splitext(bam_output)[0] + ".xml"
            ds_new = dataset_type(bam_output)
            ds_new.write(ds_out)
            xml_files.append(ds_out)
    ds_new = dataset_type(*xml_files)
    ds_new.write(output_file)
    logging.info("%s saved as %s" % (dataset_type.__name__, output_file))
    return output_file
