"""Create new columns from expanding a regexp match on an existing column.
"""
# Importing the version number (in __init__.py) and the package name
from pyaddons import (
    __version__,
    __name__ as pkg_name,
    utils,
    log
)
from pyaddons.flat_files import header
from tqdm import tqdm
import stdopen
import argparse
import sys
import os
import csv
import re
# import pprint as pp


_SCRIPT_NAME = 'regex-column'
"""The name of the script that is output when verbose is True (`str`)
"""
_DESC = __doc__
"""The program description (`str`)
"""
_DEFAULT_DELIMITER = "\t"
"""The default delimiter (`str`)
"""
_DEFAULT_COMMENT = "#"
"""The default comment character (`str`)
"""
_DEFAULT_ENCODING = "UTF-8"
"""The default encoding (`str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script.
    """
    # Initialise and parse the command line arguments
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)

    logger = log.init_logger(
        _SCRIPT_NAME, verbose=args.verbose
    )
    log.log_prog_name(logger, pkg_name, __version__)
    log.log_args(logger, args)

    prog_verbose = log.progress_verbose(args.verbose)

    # Now build the proper regex arguments
    regex_args = \
        [(i, args.regex[idx]) for idx, i in enumerate(args.column)]

    try:
        no_match = regex_columns(
            args.infile, regex_args, outfile=args.outfile,
            tmpdir=args.tmpdir, verbose=prog_verbose,
            delimiter=args.delimiter, comment=args.comment,
            encoding=args.encoding
        )
        logger.info(f"there were {no_match} groups that did not match")
        log.log_end(logger)
    except (BrokenPipeError, KeyboardInterrupt):
        # Python flushes standard streams on exit; redirect remaining
        # output to devnull to avoid another BrokenPipeError at shutdown
        devnull = os.open(os.devnull, os.O_WRONLY)
        try:
            os.dup2(devnull, sys.stdout.fileno())
        except ValueError as e:
            if not e.args[0] == "I/O operation on closed file":
                raise

        log.log_interrupt(logger)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():
    """Initialise the command line arguments and return the parser.

    Returns
    -------
    args : :obj:`argparse.ArgumentParser`
        The argparse parser object with arguments added
    """
    parser = argparse.ArgumentParser(
        description=_DESC
    )
    parser.add_argument(
        '-i', '--infile', type=str,
        help="Input files to process, if not provided input is assumed"
        " to be from STDIN."
    )
    parser.add_argument(
        '-o', '--outfile', type=str,
        help="The path to the output file (if not given will default to"
        " STDOUT), if it has a .gz extension output will be compressed"
    )
    parser.add_argument(
        '-d', '--delimiter', type=str, default="\t",
        help="The delimiter of the input (and output) file"
    )
    parser.add_argument(
        '-e', '--encoding', type=str, default="UTF8",
        help="The encoding of the input file"
    )
    parser.add_argument(
        '-c', '--comment', type=str, default="#",
        help="Skip lines that start with this string"
    )
    parser.add_argument(
        '-T', '--tmpdir', type=str, default="#",
        help="Alternative tmpdir to use, output is written to tmp before"
        " final output (unless STDOUT)"
    )
    parser.add_argument(
        '-v', '--verbose',  action="count", default=0,
        help="give more output, -vv will turn on progress monitoring"
    )
    parser.add_argument(
        "-C", "--column", nargs='+',
        help="one or more columns to apply regular expressions to. The number"
        " of ``--column`` arguments must equal the number of ``--regex``"
        " arguments"
    )
    parser.add_argument(
        "-R", "--regex", nargs='+',
        help="One or more regexps to apply to ``--column``. The number"
        " of ``--regex`` arguments must equal the number of ``--column``"
        " arguments. The named capture groups must be unique across all the"
        " regular expressions and are used for column names. These are python"
        " named capture groups, so ``(?P<NAME>regexp)`` and for no capture "
        "groups ``(?:regexp)``"
    )
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):
    """Parse the command line arguments.

    Parameters
    ----------
    parser : :obj:`argparse.ArgumentParser`
        The argparse parser object with arguments added.

    Returns
    -------
    args : :obj:`argparse.Namespace`
        The argparse namespace object containing the arguments
    """
    args = parser.parse_args()

    if len(args.column) != len(args.regex):
        raise ValueError("length of --regex must equal length of --column")

    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def regex_columns(infile, regexps, outfile=None, tmpdir=None, verbose=False,
                  delimiter=_DEFAULT_DELIMITER, comment=_DEFAULT_COMMENT,
                  encoding=_DEFAULT_ENCODING, ignorecase=False):
    """The main high level API function that takes an input file and generates
    and output file with the regex extracted columns added.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    regexps : `list` of `tuple`
        Regular expressions to match against the columns. Each tuple
        are column names and the regular expressions to apply to the
        data within them. The regular expressions should contain named capture
        groups that will be used as output column names i.e.
        ``(?P<NAME>regexp)``. The named capture groups must be unique across
        all the column/regex pairs and there must be at least 1 named capture
        group per column. Please note that this function will update this
        dictionary with a ``NO_MATCH`` field after it completes. This will have
        the number of group extractions that gave no data.
    outfile : `str` or `NoneType`, optional, default: `NoneType`
        The output file. If ``NoneType`` then output will be to ``STDOUT``.
    tmpdir : `str` or `NoneType`, optional, default: `NoneType`
        If writing to a file, a temp file is used and copied over to the final
        location when finished. Use this to specify a different temp location.
        If `NoneType` then the system default temp location is used.
    verbose : `bool`, optional, default: `False`
        Report progress.
    delimiter : `str`, optional, default `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        The symbol for a comment row that will be skipped if they occur before
        the header.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    ignorecase : `bool`, optional, default: False
        Compile the regular expressions with the `re.IGNORECASE` flag

    Returns
    -------
    no_match : `int`
        The number of group extractions that did not return any data, i.e.
        `NoneType` or ''.
    """
    no_match = dict()
    write_method = utils.get_open_method(outfile, allow_none=True)
    with stdopen.open(outfile, mode='wt', method=write_method, use_tmp=True,
                      tmpdir=tmpdir, encoding=encoding) as out:
        writer = csv.writer(out, delimiter=delimiter,
                            lineterminator=os.linesep)
        for row in tqdm(
                yield_regex_columns(infile, regexps, no_match, comment=comment,
                                    yield_header=True, delimiter=delimiter,
                                    encoding=encoding, ignorecase=ignorecase),
                desc="[progress] regex rows: ",
                unit=" rows",
                disable=not verbose
        ):
            writer.writerow(row)
    return no_match['NO_MATCH']


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_regex_columns(infile, regexps, no_match_ref, yield_header=True,
                        comment=_DEFAULT_COMMENT,
                        delimiter=_DEFAULT_DELIMITER,
                        encoding=_DEFAULT_ENCODING, ignorecase=False):
    """Yield regexp expanded rows.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    regexps : `list` of `tuple`
        Regular expressions to match against the columns. Each tuple
        are column names and the regular expressions to apply to the
        data within them. The regular expressions should contain named capture
        groups that will be used as output column names i.e.
        ``(?P<NAME>regexp)``. The named capture groups must be unique across
        all the column/regex pairs and there must be at least 1 named capture
        group per column. Please note that this function will update this
        dictionary with a ``NO_MATCH`` field after it completes. This will have
        the number of group extractions that gave no data.
    no_match_ref : `dict`
        A dictionary to place the number of no regular expression matches into.
        This is used as this function is a generator and we can't return from
        it, so we pass by reference instead.
    yield_header : `bool`, optional, default: `True`
        When processing the input file the first row yielded should be a header
        row.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        The symbol for a comment row that will be skipped if they occur before
        the header.
    delimiter : `str`, optional, default `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    ignorecase : `bool`, optional, default: False
        Compile the regular expressions with the `re.IGNORECASE` flag

    Yields
    ------
    row : `list`
        A row of data. If ``yield_header=True`` then the first row will be the
        header row.
    """
    # Parse the capture groups from the regexps and pre-compile them
    column_regex = parse_all_regex(regexps, ignorecase=ignorecase)

    # Keep track of the number of group extractions that did not match anything
    no_match = 0
    read_method = utils.get_open_method(infile, allow_none=True)

    with stdopen.open(infile, method=read_method, encoding=encoding) as incsv:
        reader = csv.reader(incsv, delimiter=delimiter)
        h, nrow = header.get_header_from_reader(reader, comment=comment)
        # Extract all the new columns that will be added from the regexp
        # extractions
        new_columns = []
        for i in column_regex:
            new_columns.extend(i[2])

            # Add the column number alongside all the regexp info
            i.append(h.index(i[0]))

        if yield_header is True:
            yield h + new_columns

        for idx, row in enumerate(reader, 1):
            # Loop through the regexps and search/add data
            for column, regex, groups, col_idx in column_regex:
                match = regex.search(row[col_idx])

                for c in groups:
                    try:
                        match_str = match.group(c)
                    except AttributeError:
                        match_str = None

                    if match_str is None or match_str == '':
                        no_match += 1
                    row.append(match_str)
                match = None

            yield row

    # A hacky way to return the no data info
    no_match_ref['NO_MATCH'] = no_match


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def parse_all_regex(column_regex, ignorecase=False):
    """Parse all the named capture groups from the column regular expressions
    and compile the regular expressions.

    Parameters
    ----------
    column_regex : `dict`
        Regular expressions to match against the columns. The keys of the dict
        are column names and the values are regular expressions to apply to the
        data within them. The regular expressions should contain named capture
        groups that will be used as output column names i.e.
        ``(?P<NAME>regexp)``. The named capture groups must be unique across
        all the column/regex pairs and there must be at least 1 named capture
        group per column.
    ignorecase : `bool`, optional, default: `False`
        Compile the regular expressions with the `re.IGNORECASE` flag.

    Returns
    -------
    parsed_regexp : `list` of `list`
        The processed regular expressions. The length of ``parsed_regexp`` will
        equal the length of the input ``column_regexp``. Each sub list has
        ``[0]`` column name, ``[1]`` compiled regexp, ``[2]`` capture groups
        (columns that will be created).

    Raises
    ------
    KeyError
        If there are repeated capture group names in the regular expression
        across the columns, or if no capture groups are identified in the
        regular expressions.
    """
    all_groups = []
    compiled = []
    # for column, regex in column_regex.items():
    for column, regex in column_regex:
        groups = _extract_groups(regex)
        all_groups.extend(groups)

        if ignorecase is False:
            compiled.append([column, re.compile(regex), groups])
        else:
            compiled.append([column, re.compile(regex, re.IGNORECASE), groups])

    if len(all_groups) > len(set(all_groups)):
        raise KeyError("capture groups not unique in regex")
    return compiled


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_groups(regex):
    """Extract all the named capture groups from a regular expression.

    Parameters
    ----------
    regex : `str`
        An un-compiled regular expression pattern.

    Returns
    -------
    capture_groups : `list` of `str`
        The named capture groups extracted from the regular expression.

    Raises
    ------
    KeyError
        If no capture groups are identified in the regular expression.
    """
    groups = []
    for i in re.findall(r'\?P<(\w+)>', regex):
        groups.append(i)

    if len(groups) == 0:
        raise KeyError(
            "no regex named capture groups found in: '{0}'".format(regex)
        )
    return groups


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == '__main__':
    main()
