"""Perform a database-like group concatenation on a file.

This is where replicated values in key columns are made unique and values in
 other columns are flattened and concatenated with a delimiter value.

IMPORTANT: The input file should be pre-sorted on the key columns.
"""
# 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 merge_sort import (
#     # chunks,
#     # heapq_merge,
#     common as com
# )
from tqdm import tqdm
from operator import itemgetter
import stdopen
import argparse
import csv
import os
import sys
# import tempfile
# import pprint as pp


_SCRIPT_NAME = 'grpcc'
"""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`)
"""

# Max sure that CSV files are as wide as possible
csv.field_size_limit(sys.maxsize)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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)

    try:
        grpcc(
            args.infile, args.keys, sort=False, distinct=args.distinct,
            not_distinct=args.not_distinct, verbose=prog_verbose,
            group_delimiter=args.group_delimiter, comment=args.comment,
            outfile=args.outfile, tmpdir=args.tmpdir,
            delimiter=args.delimiter, encoding=args.encoding
        )
        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(
        "keys", nargs='+',
        help="One or more columns to unique on."
    )
    parser.add_argument(
        '-i', '--infile', type=str,
        help="Input file to group concat, if not provided input is assumed"
        " to be from STDIN. The input file should be pre-sorted on the key"
        " columns."
    )
    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(
    #     "--sort", action="store_true",
    #     help="Pre-sort the input file i.e. the input file not already "
    #     "sorted on the keys (not implemented yet - please use unix sort and "
    #     "pipe into grpcc)"
    # )
    parser.add_argument(
        "--group-delimiter", type=str, default='|',
        help="The delimiter to use to separate the grouped values. "
        "(default is pipe '|')"
    )
    parser.add_argument(
        "--distinct", nargs='+', default=['...'], type=str,
        help="The column names where aggregated values should be made unique."
        " The default is all of them"
    )
    parser.add_argument(
        "--not-distinct", nargs='+', default=[], type=str,
        help="The column names where aggregated values should NOT be made"
        " unique. The default is all of them"
    )
    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()

    # The user has set no-distinct to all the columns to we make sure that
    # the default value for distinct is removed to avoid column collisions
    if args.not_distinct == ['...'] and args.distinct == ['...']:
        args.distict = []

    # Check the distinct arguments
    _check_distinct(args.distinct, args.not_distinct)

    return args


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _Holder(object):
    """A class to wrap around a list or a set and provide a common interface
    to store the values to be aggregated for a column.

    Parameters
    ----------
    delimiter : `str`, optional, default: `|`
        The delimiter of the aggregated data.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, delimiter='|'):
        self.delimiter = delimiter
        self.values = None

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add(self, new_value):
        """Add a value to the object.

        Parameters
        ----------
        new_value : `Any`
            The value being added to the values within the object
        """
        pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flat(self):
        """Flatten the aggregated values into a string separated by
        ``delimiter``

        Returns
        -------
        flat_str : `str`
            The flatten values stored within the object
        """
        return self.delimiter.join(
            [i if i is not None else '' for i in self.values]
        )


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _ListHolder(_Holder):
    """Implements a list-like functionality in that it only all values
    provided.

    Notes
    -----
    This results in replicated values in the flattened string.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.values = list()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add(self, new_value):
        """Add a value to the object.

        Parameters
        ----------
        new_value : `Any`
            The value being added to the values within the object
        """
        self.values.append(new_value)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _SetHolder(_Holder):
    """Implements a set-like functionality in that it only stores unique
    values.

    Notes
    -----
    Note that this uses a dict and not a set for the implementation as in
    Python >3.7 the dict will be ordered by default, so distinct values will be
    output in the order they are added.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.values = dict()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add(self, new_value):
        """Add a value to the object.

        Parameters
        ----------
        new_value : `Any`
            The value being added to the values within the object
        """
        self.values[new_value] = True

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flat(self):
        """Flatten the aggregated values into a string separated by
        ``delimiter``

        Returns
        -------
        flat_str : `str`
            The flatten values stored within the object
        """
        return self.delimiter.join(
            [i if i is not None else '' for i in self.values.keys()]
        )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def grpcc(infile, keys, sort=True, distinct=['...'], not_distinct=None,
          group_delimiter='|', outfile=None, tmpdir=None, verbose=False,
          delimiter=_DEFAULT_DELIMITER, comment=_DEFAULT_COMMENT,
          encoding=_DEFAULT_ENCODING):
    """The main high level API function that takes an input file and generates
    and output file here the values in the key columns are unique and all other
    values are aggregated.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    keys : `list` of `str`
        The key column names to make unique.
    sort : `bool`, optional, default: `True`
        Sort the input file on the key columns, if `False` it is assumed that
        it is already presorted (this is not checked).
    distinct : `list` of `str`, optional, default: `['...']`
        The column names where aggregated values will be made distinct
        (unique). If it is a list containing `...` then this is shorthand for
        all the columns
    not_distinct : `NoneType` or `list` of `str`, optional, default: `NoneType`
        The column names where aggregated values will not be made distinct
        (unique). If it is a list containing `...` then this is shorthand for
        all the columns. It is an error to have overlapping column names in
        ``distinct`` and ``not_distinct``.
    group_delimiter : `str`, optional, default: `|`
        The delimiter for the aggregated 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.
    """
    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_grpcc(infile, keys, sort=sort, distinct=distinct,
                            not_distinct=not_distinct, comment=comment,
                            group_delimiter=group_delimiter,
                            yield_header=True, delimiter=delimiter,
                            encoding=encoding, tmpdir=tmpdir),
                desc="[progress] rows: ",
                unit=" rows",
                disable=not verbose
        ):
            writer.writerow(row)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_grpcc(infile, keys, sort=True, distinct=['...'], not_distinct=None,
                group_delimiter='|', comment=_DEFAULT_COMMENT,
                yield_header=True, delimiter=_DEFAULT_DELIMITER,
                encoding=_DEFAULT_ENCODING, tmpdir=None):
    """Yield regexp expanded rows.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    keys : `list` of `str`
        The key column names to make unique.
    sort : `bool`, optional, default: `True`
        Sort the input file on the key columns, if `False` it is assumed that
        it is already presorted.
    distinct : `list` of `str`, optional, default: `['...']`
        The column names where aggregated values will be made distinct
        (unique). If it is a list containing `...` then this is shorthand for
        all the columns
    not_distinct : `NoneType` or `list` of `str`, optional, default: `NoneType`
        The column names where aggregated values will not be made distinct
        (unique). If it is a list containing `...` then this is shorthand for
        all the columns. It is an error to have overlapping column names in
        ``distinct`` and ``not_distinct``.
    group_delimiter : `str`, optional, default: `|`
        The delimiter for the aggregated data.
    comment : `str`, optional, default: `#`
        The symbol for a comment row that will be skipped if they occur before
        the header.
    yield_header : `bool`, optional, default: `True`
        When processing the input file the first row yielded should be a header
        row.
    delimiter : `str`, optional, default `\t`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `UTF-8`
        The default encoding of the input (and output) files.
    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.

    Yields
    ------
    row : `list`
        A row of data. If ``yield_header=True`` then the first row will be the
        header row.
    """
    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)

        # Get the column numbers and the data type for the sort columns
        # The data type is set to None which means whatever was input.
        key_def = [(h.index(i), None) for i in keys]

        # Get the column numbers of the index columns in the input file
        key_idx = [idx for idx, data_type in key_def]

        # Generate a function definition that will provide the keys for
        # sorting and aggregating. This will accept a row (list) and return
        # a tuple of keys
        key_func = itemgetter(*key_idx)

        # This will hold the column numbers of all the non-index columns
        # that will hold the aggregated data.
        update_idx = list(range(len(h)))
        # Reverse sort and remove the index columns
        for i in sorted(key_idx, reverse=True):
            update_idx.pop(i)

        # Used to flatten the input, this are class definitions to hold the
        # aggregated data. Either list holders or set holders.
        distinct_func = _set_distinct(h, distinct, not_distinct)

        # Will hold the current key being processed and the current
        # aggregations for the key
        cur_key = None
        holder_class = []

        if yield_header is True:
            yield h

        if sort is True:
            raise NotImplementedError("please pre-sort the input file")
        else:
            # Assume the input is pre-sorted
            for row in reader:
                # Get the keys for the row
                key = key_func(row)
                # print(key, file=sys.stderr)
                # Is there a change in keys? If so that is the indicator
                # that we need to output the aggregated row
                if key != cur_key:
                    try:
                        yield _build_output_row(
                            update_idx, key_idx, cur_key, holder_class
                        )
                    except (IndexError, TypeError):
                        # This is probably the first iteration and cur_key
                        # is None, but we check just in case
                        if cur_key is not None:
                            raise
                    # Update the key for the new key that we have
                    cur_key = key

                    # Initialise the new aggregators for the new key
                    holder_class = \
                        [i(delimiter=group_delimiter) for i in distinct_func]

                # Add the contents of the current row to the aggregators
                for idx, i in enumerate(holder_class):
                    i.add(row[idx])

            # After parsing we will still have one last key to output
            yield _build_output_row(
                update_idx, key_idx, cur_key, holder_class
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _build_output_row(update_idx, key_idx, cur_key, holders):
    """Assemble an output row from the aggregations and the current key.

    Parameters
    ----------
    update_idx : `list` of `int`
        The column numbers of the non-index columns, the matching holders at
        these locations  will have the aggregated data for the row.
    key_idx : `list` of `int`
        The column numbers of the key columns, this will be used to update the
        key (unique non-aggregated) data for the row, this actual values are
        stored in ``cur_key`` and match the order of the ``key_idx``.
    cur_key : `tuple` of `str`
        The data values for the current key that will be placed in the
        assembled row.
    holders : `list` of (`pyaddons.flat_files.grpcc._ListHolder` or `pyaddons.flat_files.grpcc._SetHolder`)
        The aggregated data holders, that will provide delimited strings for
        the assembled row.

    Returns
    -------
    row : `list` of `str`
        The row of aggregated data.
    """
    row = [None] * len(holders)

    if not isinstance(cur_key, tuple):
        cur_key = cur_key,
    # print(key_idx, file=sys.stderr)
    # print(cur_key, file=sys.stderr)
    for idx, value in zip(key_idx, cur_key):
        row[idx] = value

    for idx in update_idx:
        row[idx] = holders[idx].flat()
    return row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _set_distinct(h, distinct, not_distinct):
    """Initialise the type of holder classes that will be used to provide the
    aggregated data for a key.

    Parameters
    ----------
    header : `list` of `str`
        The header of the input file.
    distinct : `list` of `str`
        The column names of the columns where we want distinct aggregations. If
        this contains a single element of ``...``, then all the header columns
        will be used.
    not_distinct : `list` of `str`
        The column names of the columns where we do not want distinct
        aggregations. If this contains a single element of ``...``, then all
        the header columns will be used.

    Returns
    -------
    holders : `list` of `classes`
        These are class definitions of the type: `bio_misc.grpcc._ListHolder`
        or `bio_misc.grpcc._SetHolder`, depending on if the column was defined
        as not distinct or distinct respectively.
    """
    if distinct == ['...']:
        distinct = h.copy()
    elif not_distinct == ['...']:
        not_distinct = h.copy()

    for i in not_distinct:
        try:
            distinct.pop(distinct.index(i))
        except ValueError:
            pass

    distinct_func = []
    for i in h:
        if i in distinct:
            distinct_func.append(_SetHolder)
        elif i in not_distinct:
            distinct_func.append(_ListHolder)
        else:
            raise ValueError("unknown column: {0}".format(i))
    return distinct_func


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _check_distinct(distinct, not_distinct):
    """Check that the distinct and not distinct definitions are valid.

    Parameters
    ----------
    distinct : `list` of `str`
        The column names of the columns where we want distinct aggregations. If
        this contains a single element of ``...``, then all the header columns
        will be used.
    not_distinct : `list` of `str`
        The column names of the columns where we do not want distinct
        aggregations. If this contains a single element of ``...``, then all
        the header columns will be used.

    Raises
    ------
    ValueError
        If both ``distinct`` and ``not_distinct`` are ``['...']`` (meaning all
        columns).
    KeyError
        If any individual columns defined in ``distinct`` and ``not_distinct``
        overlap.
    """
    if distinct == ['...'] and not_distinct == ['...']:
        raise ValueError(
            "both distinct and not distinct are set to all columns"
        )
    # Now make sure there are no overlaps between no_distinct and distinct
    overlaps = set(not_distinct).intersection(distinct)

    if len(overlaps) > 0:
        raise KeyError(
            "overlapping distinct and no-distinct column names: "
            "{0}".format(",".join(overlaps))
        )


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