from __future__ import division, unicode_literals, absolute_import

from builtins import int, range, dict, map, zip

import io
import os
import re
import sys
import queue

# Future warning from cython in h5py
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)
import h5py

from tqdm import tqdm
from time import sleep
from itertools import islice
from multiprocessing import Process, Queue, Pipe

if sys.version_info[0] > 2:
    unicode = str

from . import tombo_helper as th


VERBOSE = False

_MAX_QUEUE_SIZE = 1000
_ITER_QUEUE_LIMIT = 1000
_PROC_UPDATE_INTERVAL = 100

_MAX_FASTQ_QUEUE_SIZE = 10000
_SEQ_SUMMARY_FN_FIELD = 'filename'
_SEQ_SUMMARY_ID_FIELD = 'read_id'

# warning messages for annotate with fastqs over multiple processes,
# requiring passing warning codes to only print warning once.
_WARN_ID_VAL = 'ids'
_WARN_IO_VAL = 'io'
_WARN_MISMATCH_VAL = 'mismatch'
_WARN_OVRWRT_VAL = 'overwrite'
_WARN_UNIQ_VAL = 'uniq'
_WARN_CODES = (_WARN_ID_VAL, _WARN_IO_VAL, _WARN_MISMATCH_VAL, _WARN_OVRWRT_VAL)
_WARN_CODES_PREP = (_WARN_OVRWRT_VAL, _WARN_UNIQ_VAL)
_WARN_PREFIX = '****** WARNING ****** '


##########################
###### Annotate Raw ######
##########################

def _prep_fast5_for_fastq(fast5_data, bc_grp_name, bc_subgrp_name, overwrite):
    read_id = th.get_raw_read_slot(fast5_data).attrs.get('read_id')
    try:
        read_id = read_id.decode()
    except (AttributeError, TypeError):
        pass
    if read_id is None:
        return

    # if Analyses group doesn't exist yet, create it
    try:
        analyses_grp = fast5_data['/Analyses']
    except:
        analyses_grp = fast5_data.create_group('Analyses')

    # create Fastq slot, unless value exists and --overwrite is not set
    try:
        bc_grp = analyses_grp[bc_grp_name]
        bc_subgrp = analyses_grp[bc_subgrp_name]
    except:
        try:
            bc_grp = analyses_grp.create_group(bc_grp_name)
            bc_subgrp = bc_grp.create_group(bc_subgrp_name)
        except:
            if overwrite:
                del analyses_grp[bc_grp_name]
                bc_grp = analyses_grp.create_group(bc_grp_name)
                bc_subgrp = bc_grp.create_group(bc_subgrp_name)
            else:
                raise th.TomboError(
                    bc_grp_name + ' exists and --overwrite is not set.')

    return read_id

def _annotate_with_fastqs_worker(
        fastq_rec_q, fast5s_read_ids, fastq_slot, fq_slot_prepped,
        prog_q, warn_q, bc_grp_name, bc_subgrp_name, overwrite):
    been_warned = dict((warn_code, False) for warn_code in _WARN_CODES)
    num_recs_proc = 0
    while True:
        fastq_rec = fastq_rec_q.get()
        if fastq_rec is None:
            break

        # extract read_id from fastq (which should be the first text after
        # the "@" record delimiter up to the first white space or underscore
        read_id = fastq_rec[0].split()[0].split('_')[0][1:]
        if read_id not in fast5s_read_ids:
            if not been_warned[_WARN_ID_VAL]:
                been_warned[_WARN_ID_VAL] = True
                warn_q.put(_WARN_ID_VAL)
            continue

        try:
            with h5py.File(fast5s_read_ids[read_id], 'r+') as fast5_data:
                if not fq_slot_prepped:
                    try:
                        file_parsed_id = _prep_fast5_for_fastq(
                            fast5_data, bc_grp_name, bc_subgrp_name, overwrite)
                    except th.TomboError:
                        if not been_warned[_WARN_OVRWRT_VAL]:
                            been_warned[_WARN_OVRWRT_VAL] = True
                            warn_q.put(_WARN_OVRWRT_VAL)
                        continue
                    if read_id != file_parsed_id:
                        if not been_warned[_WARN_MISMATCH_VAL]:
                            been_warned[_WARN_MISMATCH_VAL] = True
                            warn_q.put(_WARN_MISMATCH_VAL)
                        continue
                bc_slot = fast5_data[fastq_slot]
                # add sequence to fastq slot
                bc_slot.create_dataset(
                    'Fastq', data=''.join(fastq_rec),
                    dtype=h5py.special_dtype(vlen=unicode))

                # progress q update
                num_recs_proc += 1
                if num_recs_proc % _PROC_UPDATE_INTERVAL == 0:
                    prog_q.put(_PROC_UPDATE_INTERVAL)
        except:
            if not been_warned[_WARN_IO_VAL]:
                been_warned[_WARN_IO_VAL] = True
                warn_q.put(_WARN_IO_VAL)
            continue

    # add last number of records reported from this process
    prog_q.put(num_recs_proc % _PROC_UPDATE_INTERVAL)

    return

def _feed_seq_records_worker(fastq_fns, fastq_rec_q, num_processes):
    for fastq_fn in fastq_fns:
        n_recs = 0
        with io.open(fastq_fn) as fastq_fp:
            while True:
                fastq_rec = list(islice(fastq_fp, 4))
                # if record contains fewer than 4 lines this indicates the
                # EOF, so move to next file
                if len(fastq_rec) != 4: break
                # if sequence identifier line does not start with "@" or quality
                # score line does not start with a "+" the file may be
                # corrupted, so don't process any more records
                if (re.match('@', fastq_rec[0]) is None or
                    re.match('\+', fastq_rec[2]) is None):
                    # TODO maybe send this as a warning code to avoid poorly
                    # formatted output
                    th.warning_message(
                        'Successfully parsed ' + unicode(n_recs) +
                        ' FASTQ records from ' + fastq_fn + ' before ' +
                        'encountering an invalid record. The rest of ' +
                        'this file will not be processed.')
                    break
                n_recs += 1
                fastq_rec_q.put(fastq_rec)

    # put none records to trigger annotation processes to exit
    for _ in range(num_processes):
        fastq_rec_q.put(None)

    return

def _get_ann_queues(prog_q, warn_q, num_read_ids, wp_conn):
    if VERBOSE: bar = tqdm(total=num_read_ids, smoothing=0)
    been_warned = dict((warn_code, False) for warn_code in _WARN_CODES)

    def update_warn(warn_val):
        if warn_val == _WARN_ID_VAL:
            if VERBOSE and not been_warned[_WARN_ID_VAL]:
                bar.write(
                    _WARN_PREFIX + 'Some FASTQ records contain read ' +
                    'identifiers not found in any FAST5 files or ' +
                    'sequencing summary files.',
                    file=sys.stderr)
            been_warned[_WARN_ID_VAL] = True
        elif warn_val == _WARN_IO_VAL:
            if VERBOSE and not been_warned[_WARN_IO_VAL]:
                bar.write(
                    _WARN_PREFIX + 'Some read files could not be accessed.',
                    file=sys.stderr)
            been_warned[_WARN_IO_VAL] = True
        elif warn_val == _WARN_MISMATCH_VAL:
            if VERBOSE and not been_warned[_WARN_MISMATCH_VAL]:
                bar.write(
                    _WARN_PREFIX + 'Read ID(s) found in sequencing summary ' +
                    'and FAST5 file are discordant. Skipping such reads.',
                    file=sys.stderr)
            been_warned[_WARN_MISMATCH_VAL] = True
        elif warn_val == _WARN_OVRWRT_VAL:
            if VERBOSE and not been_warned[_WARN_OVRWRT_VAL]:
                bar.write(
                    _WARN_PREFIX + 'Basecalls exsit in specified slot for ' +
                    'some reads. Set --overwrite option to overwrite these ' +
                    'basecalls.', file=sys.stderr)
            been_warned[_WARN_OVRWRT_VAL] = True
        else:
            if VERBOSE:
                bar.write('{}Invalid warning code encountered: {}'.format(
                    _WARN_PREFIX, warn_val), file=sys.stderr)

        return


    total_added_seqs = 0
    while True:
        try:
            iter_added = prog_q.get(block=False)
            total_added_seqs += iter_added
            if VERBOSE: bar.update(iter_added)
        except queue.Empty:
            try:
                warn_val = warn_q.get(block=False)
                update_warn(warn_val)
            except queue.Empty:
                sleep(0.1)
                # check if main thread has finished with all fastq records
                if wp_conn.poll():
                    break

    # collect all remaining warn and progress values
    while not prog_q.empty():
        iter_added = prog_q.get(block=False)
        total_added_seqs += iter_added
        if VERBOSE: bar.update(iter_added)
    while not warn_q.empty():
        warn_val = warn_q.get(block=False)
        update_warn(warn_val)

    if VERBOSE:
        bar.close()
        th.status_message('Added sequences to a total of ' +
                          str(total_added_seqs) + ' reads.')
        if total_added_seqs < num_read_ids:
            th.warning_message(
                'Not all read ids from FAST5s or sequencing summary files ' +
                'were found in FASTQs.\n\t\tThis can result from reads that ' +
                'failed basecalling or if full sets of FAST5s/sequence ' +
                'summaries are not processed with full sets of FASTQs.')

    return

def _annotate_with_fastqs(
        fastq_fns, fast5s_read_ids, fastq_slot, fq_slot_prepped, num_processes,
        bc_grp_name, bc_subgrp_name, overwrite):
    if VERBOSE: th.status_message('Annotating FAST5s with sequence from FASTQs.')
    fastq_rec_q = Queue(maxsize=_MAX_FASTQ_QUEUE_SIZE)
    prog_q = Queue()
    warn_q = Queue()

    # open a single process to read fastq files and feed the fastq record queue
    fq_feed_p = Process(target=_feed_seq_records_worker,
                        args=(fastq_fns, fastq_rec_q, num_processes))
    fq_feed_p.daemon = True
    fq_feed_p.start()

    # open fast5 annotation processes
    ann_args = (fastq_rec_q, fast5s_read_ids, fastq_slot, fq_slot_prepped,
                prog_q, warn_q, bc_grp_name, bc_subgrp_name, overwrite)
    ann_ps = []
    for p_id in range(num_processes):
        ann_p = Process(target=_annotate_with_fastqs_worker, args=ann_args)
        ann_p.daemon = True
        ann_p.start()
        ann_ps.append(ann_p)

    main_wp_conn, wp_conn = Pipe()
    warn_prog_p = Process(target=_get_ann_queues,
                          args=(prog_q, warn_q, len(fast5s_read_ids), wp_conn))
    warn_prog_p.daemon = True
    warn_prog_p.start()

    fq_feed_p.join()
    for ann_p in ann_ps:
        ann_p.join()
    # send signal to warn/progress queue that all other processes are complete
    main_wp_conn.send(True)
    warn_prog_p.join()

    return


##########################
#### Extract read_ids ####
##########################

def _get_prep_queue(read_ids_q, prog_q, warn_q, gp_conn, num_fast5s):
    """Process all records from all fast5 prep queues
    """
    ovrwrt_mess = (
        _WARN_PREFIX + 'Basecalls exsit in specified slot for some ' +
        'reads. Set --overwrite option to overwrite these basecalls.')
    fast5s_read_ids = {}
    # Warn non-unique read_ids in directory
    been_warned = dict((warn_code, False) for warn_code in _WARN_CODES_PREP)
    if VERBOSE: bar = tqdm(total=num_fast5s, smoothing=0)

    while True:
        try:
            read_id, fast5_fn = read_ids_q.get(block=False)
            if read_id in fast5s_read_ids:
                if VERBOSE and not been_warned[_WARN_UNIQ_VAL]:
                    bar.write(
                        _WARN_PREFIX + 'Multiple FAST5 files contain the ' +
                        'same read ID. Ensure that FAST5 files are from a ' +
                        'single run.', file=sys.stderr)
                    been_warned[_WARN_UNIQ_VAL] = True
                continue
            fast5s_read_ids[read_id] = fast5_fn
        except queue.Empty:
            try:
                warn_val = warn_q.get(block=False)
                if warn_val == _WARN_OVRWRT_VAL:
                    if VERBOSE and not been_warned[_WARN_OVRWRT_VAL]:
                        bar.write(ovrwrt_mess, file=sys.stderr)
                    been_warned[_WARN_OVRWRT_VAL] = True
                else:
                    bar.write('{}Invalid warning code encountered: {}'.format(
                        _WARN_PREFIX, warn_val), file=sys.stderr)
            except queue.Empty:
                try:
                    if VERBOSE: bar.update(prog_q.get(block=False))
                except queue.Empty:
                    sleep(0.1)
                    # check if main thread has finished with all FAST5s
                    if gp_conn.poll():
                        break

    while not read_ids_q.empty():
        read_id, fast5_fn = read_ids_q.get(block=False)
        fast5s_read_ids[read_id] = fast5_fn
    while not warn_q.empty():
        warn_val = warn_q.get(block=False)
        if warn_val == _WARN_OVRWRT_VAL:
            if VERBOSE and not been_warned[_WARN_OVRWRT_VAL]:
                bar.write(ovrwrt_mess, file=sys.stderr)
                been_warned[_WARN_OVRWRT_VAL] = True
        else:
            bar.write('{}Invalid warning code encountered: {}'.format(
                _WARN_PREFIX, warn_val), file=sys.stderr)
    while not prog_q.empty():
        if VERBOSE: bar.update(prog_q.get(block=False))

    if VERBOSE: bar.close()
    gp_conn.send(fast5s_read_ids)

    return

def _prep_fastq_slot_worker(
        fast5_q, bc_grp, bc_subgrp, overwrite, read_ids_q, prog_q, warn_q):
    num_files_proc = 0
    been_warned_overwrite = False
    while True:
        try:
            fast5_fn = fast5_q.get(block=False)
        except queue.Empty:
            sleep(0.1)
            continue

        if fast5_fn is None:
            break

        num_files_proc += 1
        if num_files_proc % _PROC_UPDATE_INTERVAL == 0:
            prog_q.put(_PROC_UPDATE_INTERVAL)

        try:
            with h5py.File(fast5_fn) as fast5_data:
                try:
                    read_id = _prep_fast5_for_fastq(
                        fast5_data, bc_grp, bc_subgrp, overwrite)
                except th.TomboError:
                    # avoid the warn queue getting too large by sending overwite
                    # warnings for each read from each thread
                    if not been_warned_overwrite:
                        been_warned_overwrite = True
                        warn_q.put(_WARN_OVRWRT_VAL)
                    continue
        except:
            continue
        if read_id is None:
            continue

        read_ids_q.put((read_id, fast5_fn))

    prog_q.put(num_files_proc % _PROC_UPDATE_INTERVAL)

    return

def _fill_files_queue(fast5_q, fast5_fns, num_ps):
    for fast5_fn in fast5_fns:
        fast5_q.put(fast5_fn)
    for _ in range(num_ps):
        fast5_q.put(None)

    return

def _get_read_ids_and_prep_fastq_slot(
        fast5s_dir, bc_grp, bc_subgrp, overwrite, num_processes):
    """Extract read id from /Raw group and prep fastq slots for annotation with
    associated FASTQ files.
    """
    if VERBOSE: th.status_message(
            'Preparing reads and extracting read identifiers.')
    fast5_q = Queue(maxsize=_MAX_QUEUE_SIZE)
    read_ids_q = Queue()
    prog_q = Queue()
    warn_q = Queue()

    fast5_fns = th.get_files_list(fast5s_dir)
    files_p = Process(target=_fill_files_queue,
                      args=(fast5_q, fast5_fns, num_processes))
    files_p.daemon = True
    files_p.start()

    prep_args = (fast5_q, bc_grp, bc_subgrp, overwrite, read_ids_q,
                 prog_q, warn_q)
    prep_ps = []
    for p_id in range(num_processes):
        prep_p = Process(target=_prep_fastq_slot_worker, args=prep_args)
        prep_p.daemon = True
        prep_p.start()
        prep_ps.append(prep_p)

    main_gp_conn, gp_conn = Pipe()
    get_prep_p = Process(
        target=_get_prep_queue,
        args=(read_ids_q, prog_q, warn_q, gp_conn, len(fast5_fns)))
    get_prep_p.daemon = True
    get_prep_p.start()

    # join all processes into the main thread
    files_p.join()
    for prep_p in prep_ps:
        prep_p.join()
    # send signal to get_prep queue that all other processes are complete
    main_gp_conn.send(True)
    fast5s_read_ids = main_gp_conn.recv()

    return fast5s_read_ids

def _parse_sequencing_summary_files(fast5s_dir, seq_summary_fns):
    if VERBOSE: th.status_message('Getting read filenames.')
    full_fast5_fns = {}
    # walk through directory structure searching for fast5 files
    for root, _, fns in os.walk(fast5s_dir):
        for fn in fns:
            if not fn.endswith('.fast5'): continue
            full_fast5_fns[fn] = os.path.join(root, fn)

    if VERBOSE: th.status_message('Parsing sequencing summary files.')
    fast5s_read_ids = {}
    been_warned = False
    for seq_summary_fn in seq_summary_fns:
        with open(seq_summary_fn) as fp:
            try:
                header_fields = fp.readline().split()
                fn_field = next(i for i, h_field in enumerate(header_fields)
                                if re.match(_SEQ_SUMMARY_FN_FIELD, h_field))
                id_field = next(i for i, h_field in enumerate(header_fields)
                                if re.match(_SEQ_SUMMARY_ID_FIELD, h_field))
            except:
                th.warning_message(
                    'Could not extract header information for sequencing ' +
                    'summary file: ' + seq_summary_fn)
                continue
            try:
                for line in fp:
                    rec_fields = line.split()
                    rec_short_fn = rec_fields[fn_field]
                    try:
                        rec_full_fn = full_fast5_fns[rec_short_fn]
                    except KeyError:
                        if not been_warned:
                            th.warning_message(
                                'Some FASTQ records from sequencing summaries ' +
                                'do not appear to have a matching file.')
                        been_warned = True
                        continue
                    # convert filename to full filename and link to read id
                    fast5s_read_ids[rec_fields[id_field]] = rec_full_fn
            except:
                th.warning_message(
                    'Error parsing records for sequencing ' +
                    'summary file: ' + seq_summary_fn)

    return fast5s_read_ids


##################################
###### Annotate FAST5s Main ######
##################################

def annotate_reads_with_fastq_main(args):
    global VERBOSE
    VERBOSE = not args.quiet
    th.VERBOSE = VERBOSE

    fast5s_basedir = (
        args.fast5_basedir if args.fast5_basedir.endswith('/') else
        args.fast5_basedir + '/')
    if args.sequencing_summary_filenames:
        fast5s_read_ids = _parse_sequencing_summary_files(
            fast5s_basedir, args.sequencing_summary_filenames)
        fq_slot_prepped = False
    else:
        fast5s_read_ids = _get_read_ids_and_prep_fastq_slot(
            fast5s_basedir, args.basecall_group, args.basecall_subgroup,
            args.overwrite, args.processes)
        fq_slot_prepped = True
    fastq_slot = '/'.join(('/Analyses', args.basecall_group,
                           args.basecall_subgroup))
    _annotate_with_fastqs(
        args.fastq_filenames, fast5s_read_ids, fastq_slot, fq_slot_prepped,
        args.processes, args.basecall_group, args.basecall_subgroup,
        args.overwrite)

    return


if __name__ == '__main__':
    sys.stderr.write('This is a module. See commands with `tombo -h`')
    sys.exit(1)
