"""A utility to add/remove and move columns in a flat file.

Please be aware this is not fully tested and some of the transformations may
not function as expected.

In the description below the ``[]`` surrounding the terms are not required they
are for documentation purposes only.

The idea behind this was to provide a safer alternative to using cut/paste or
awk for modifying flat files that contain a header as all references to columns
are based on column names and not on numbers, so errors should become more
apparent. Also, multiple changes can be done in a single pass through the file.

The user supplies some specific transformations to the input. These are applied
in the order that they are given.

Columns can be moved by either specifying the column name in a defined
positional order or by making the movement relative to another column (the
anchor column) by using the format ``[column to move]<[anchor]``, to move
``[column to move]`` before the ``[anchor]`` column. A > can be used to move
the column after the anchor column.

Columns can also be renamed and moved at the same time by specifying
``[old column name]=[new column name]``, in a defined positional order.

If you want to rename the column but keep it in the same position (relative to
everything else that is happening around it), then you can use the format.
``=[old column name]=[new column name]``. Where the column eventually ends up
will depend on the other transformations.

New columns can be added by using the format ``+[new column name]`` for an
empty column or ``+[new column name]=[initial value]``. With this format the
position of the new column will be in the positional sequence in which it has
been defined. It is possible to create a new column relative to another column
using the format ``+[new column name]=[initial value]<[anchor column]`` or
``+[new column name]=[initial value]>[anchor column]``. In this context
[initial value] can be blank.

Columns can be specifically removed using the format ``-[column name]``. Note
that the default is to remove any column not explicitly specified in a
transform, however, if ``keep_existing=True``, then columns that are not
included in a transform are kept in place relative to all the transforms around
them. In this context ``-[column name]`` may be useful.
"""
# 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 argparse
import stdopen
import sys
import os
import csv
import re
# import pprint as pp


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

# The locations of attributes in the parsed transforms
_OLD_COL = 0
"""The element index location of the old column attribute (`int`)
"""
_NEW_COL = 1
"""The element index location of the new column attribute (`int`)
"""
_TO_KEEP = 2
"""The element index location of the to keep attribute set in all columns
influenced by a transform (`int`)
"""
_DEFAULT_VALUE = 3
"""The element index location of the default value for a new column
attribute (`int`)
"""
_IS_NEW_COL = 4
"""The element index location of the is a new column flag attribute (`int`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script. For API access see:
    ``pyaddons.flat_files.rejig.rejig``
    """
    # 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)

    for i in args.transforms:
        logger.info(f" [transform] {i}")

    try:
        rejig(args.infile, args.transforms, outfile=args.outfile,
              tmpdir=args.tmpdir, verbose=args.verbose,
              delimiter=args.delimiter, comment=args.comment,
              skiplines=args.skip_lines, encoding=args.encoding,
              keep_existing=args.keep)
        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(
        'transforms', type=str, nargs='+',
        help="The transformations to apply."
    )
    parser.add_argument(
        '-i', '--infile', type=str,
        help="The path to the input file (can be gzip compressed)"
    )
    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(
        '-s', '--skip-lines', type=int, default=0,
        help="The skip this many lines before processing the header"
    )
    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(
        '--keep', action="store_true",
        help="Retain columns not in the transform."
    )

    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()
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def rejig(infile, transforms, outfile=None, tmpdir=None, verbose=False,
          delimiter=_DEFAULT_DELIMITER, comment=_DEFAULT_COMMENT,
          skiplines=0, encoding=_DEFAULT_ENCODING, keep_existing=False):
    """The main high level API function that takes an input file and generates
    and output file with all the transforms applied.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    transforms : `list` of `str`
        All the transforms that will be applied to the file in the order that
        you want them applied. See module level notes for details for the
        formats.
    outfile : `str` or `NoneType`, optional, default: `NoneType`
        The output file. If ``NoneType`` then output will be to ``STDOUT``.
    tmp_dir : `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.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    keep_existing : `bool`, optional, default: `False`
        Should columns that have not specifically been transformed be kept in
        the output.
    """
    open_method = open
    try:
        open_method = utils.get_open_method(
            outfile
        )
    except TypeError:
        pass

    with stdopen.open(outfile, method=open_method, mode='wt', use_tmp=True,
                      tmpdir=tmpdir, encoding=encoding) as out:
        writer = csv.writer(out, delimiter=delimiter,
                            lineterminator=os.linesep)
        for row in tqdm(
                yield_rejig(infile, transforms, comment=comment,
                            yield_header=True, delimiter=delimiter,
                            skiplines=skiplines, encoding=encoding,
                            keep_existing=keep_existing),
                desc="[progress] rejig(s): ",
                unit=" rejig(s)",
                disable=not verbose
        ):
            writer.writerow(row)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_rejig(infile, transforms, comment=_DEFAULT_COMMENT,
                yield_header=True, delimiter=_DEFAULT_DELIMITER,
                skiplines=0, encoding=_DEFAULT_ENCODING, keep_existing=False):
    """Yield transformed rows.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    transforms : `list` of `str`
        All the transforms that will be applied to the file in the order that
        you want them applied. See module level notes for details for the
        formats.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        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 `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    keep_existing : `bool`, optional, default: `False`
        Should columns that have not specifically been transformed be kept in
        the output.

    Yields
    ------
    row : `list`
        A row of data. If ``yield_header=True`` then the first row will be the
        header row.
    """
    open_method = open
    try:
        open_method = utils.get_open_method(
            infile
        )
    except TypeError:
        pass

    with stdopen.open(infile, method=open_method, encoding=encoding) as infile:
        reader = csv.reader(infile, delimiter=delimiter)
        h, row_no = header.get_header_from_reader(reader, comment=comment,
                                                  skiplines=skiplines)
        new_header, remap, template_row = parse_transforms(
            h, transforms, keep_existing=keep_existing
        )
        if yield_header is True:
            yield new_header

        for idx, row in enumerate(reader, 1):
            new_row = template_row.copy()
            for old_idx, new_idx in remap:
                new_row[new_idx] = row[old_idx]
            yield new_row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def parse_transforms(header, transforms, keep_existing=False):
    """Process the user defined transforms into re-mappings and a new header
    and a template row of default values for all columns.

    Parameters
    ----------
    header : `list` of `str`
        The existing header of the input file
    transforms : `list` of `str`
        All the transforms that will be applied to the file in the order that
        you want them applied. See module level notes for details for the
        formats.
    keep_existing : `bool`, optional, default: `False`
        Should columns that have not specifically been transformed be kept in
        the output.

    Returns
    -------

    Raises
    ------
    KeyError
        If there are duplicated column names in the existing input file header.
    ValueError
        If any of the columns in the transform are not in the header
    """
    # Define the current state of the file
    # old_col, new_col, to_keep, default_value, is_new
    output = [[i, i, False, None, False] for idx, i in enumerate(header)]

    # check that there are no duplication's in the header
    if len(header) > len(set(header)):
        raise KeyError("there are duplicated column names in the header")

    # Loop through all the user defined transforms
    # TODO: This is a bit messy and could do with cleaning up
    for idx, trans in enumerate(transforms):
        i = trans.strip()
        # Adding a column, could either be addition with no default value
        # or with a default value or addition relative to existing column
        if trans.startswith('+'):
            # TODO: Check that new/renamed columns do not create duplicates
            # Note the passed value has leading + removed
            new_col, default_value, is_after, anchor_col = \
                _process_add_column(i[1:])

            # This is the column that will e inserted, now determine where.
            insert_column = [new_col, new_col, True, default_value, True]

            # Relative insert
            if anchor_col is not None:
                anchor_idx, anchor = _find_in_output(anchor_col, output)
                output.insert(anchor_idx+is_after, insert_column)
            else:
                # TODO: This should be _find_last_to_keep(output)
                output.append(insert_column)
        elif trans.startswith('-'):
            # Remove from output
            anchor_idx, anchor = _find_in_output(trans[1:], output)
            output.pop(anchor_idx)
        elif trans.startswith('='):
            # In place rename
            old_col, new_col = _process_equals(i[1:])
            source_idx, source = _find_in_output(old_col, output)
            source[_NEW_COL] = new_col
            source[_TO_KEEP] = True
        else:
            # Moving or moving relative or move and rename
            old_col, new_col, is_after, anchor_col = \
                _process_positional_column(i)
            # Moving relative to anchor column
            if anchor_col is not None:
                anchor_idx, anchor = _find_in_output(anchor_col, output)
                source_idx, source = _find_in_output(old_col, output)
                output.pop(source_idx)
                source[_TO_KEEP] = True
                output.insert(anchor_idx+is_after, source)
            else:
                # Moving or moving and renaming
                source_idx, source = _find_in_output(old_col, output)
                output.pop(source_idx)
                source[_NEW_COL] = new_col
                source[_TO_KEEP] = True
                insert_after = _find_last_to_keep(output)
                output.insert(insert_after+1, source)

    # If we are not keeping existing columns, then remove them from output
    if keep_existing is False:
        output = [i for i in output if i[_TO_KEEP] is True]

    # Create a template row
    template_row = [i[_DEFAULT_VALUE] for i in output]

    # Create the mappings between old column locations and new ones
    remap = []
    for new_idx, i in enumerate(output):
        try:
            old_idx = header.index(i[_OLD_COL])
            remap.append((old_idx, new_idx))
        except ValueError:
            # Not in header so could be a new row, if not then we have an error
            # not sure but I think these will be caught elsewhere
            if i[_IS_NEW_COL] is False:
                raise

    return [i[_NEW_COL] for i in output], remap, template_row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _find_last_to_keep(output):
    """Get the element index of the last column that has been involved in a
    transform, this gives us the position of the next "positional" transform
    (i.e not relative).

    Parameters
    ----------
    output : `list` of `list`
        The current output list, each sub-list has ``[0]`` old column name
        ``[1]`` new column name ``[2]`` to keep flag ``[3]`` default value
        ``[4]`` is new flag

    Returns
    -------
    last_to_keep : `int`
        The index of the last column with the to keep flag set to ``True``.
        If there are no columns with the to keep flag set to ``True`` then -1
        is returned.
    """
    last = -1
    for idx, i in enumerate(output):
        if i[_TO_KEEP] is True:
            last = idx
    return last


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _find_in_output(column, output):
    """Locate a column within the current output list.

    Parameters
    ----------
    column : `str`
        The column to locate.
    output : `list` of `list`
        The current output list, each sub-list has ``[0]`` old column name
        ``[1]`` new column name ``[2]`` to keep flag ``[3]`` default value
        ``[4]`` is new flag

    Returns
    -------
    idx : `int`
        The index that the column was found at in ``output``
    column : ``list``
        The column data that matches ``column``, has the following structure:
        ``[0]`` old column name ``[1]`` new column name ``[2]`` to keep flag
        ``[3]`` default value ``[4]`` is new flag.

    Raises
    ------
    ValueError
        If the ``column`` can't be found in ``output``
    """
    for idx, i in enumerate(output):
        if i[_OLD_COL] == column:
            return idx, i

    raise ValueError(
        "can't find column in output, have you deleted/renamed it? "
        "{0}".format(column)
    )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _process_positional_column(value):
    """Process transforms that do not start with +/-/=.

    Parameters
    ----------
    value : `str`
        A transform that had no leading sign on it.

    Returns
    -------
    old_col : `str`
        The old column name.
    new_col : `str`
        The new column name.
    is_after : `bool` or `NoneType`
        If the column is being moved from a different location then this will
        be a `bool` indicating if it is after the anchor_col (`True`) or before
        it (`False`). If the column is being renamed or not moved relatively
        then this will be `NoneType`.
    anchor_col : `str` or `NoneType`
        If being moved relatively, then this is the name of the column it is
        being moved relative to. If not a relative move then this will be
        `NoneType`.

    Notes
    -----
    A positional column could be moving an existing column to a defined
    position or moving an existing column relative to another existing column
    or renaming a column at a defined position.
    """
    try:
        # First are we renaming? i.e. old_col=new_col
        old_col, new_col = _process_equals(value)
        # old, new, is_after, anchor
        return old_col, new_col, None, None
    except ValueError as e:
        # No equals
        if 'unpack' not in e.args[0]:
            raise
        # Moving or moving relative to another column

    try:
        # Is there a relative sign?
        source_col, is_after, anchor_col = _process_gt_lt(value)
        return source_col, source_col, is_after, anchor_col
    except ValueError as e:
        # No relative sign this must be just a single column name
        if 'unpack' not in e.args[0]:
            raise
    return value, value, None, None


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _process_add_column(value):
    """Process transforms that start with a plus.

    Parameters
    ----------
    value : `str`
        A transform that had a leading + sign on it (which was removed before
        passing to this function).

    Returns
    -------
    new_col : `str`
        The new column name.
    default_value : `str` of `NoneType`
        The default value for the new column, if not defined in the transform
        then it will be `NoneType`.
    is_after : `bool` or `NoneType`
        If the column is being moved from a different location then this will
        be a `bool` indicating if it is after the anchor_col (`True`) or before
        it (`False`). If the column is being renamed or not moved relatively
        then this will be `NoneType`.
    anchor_col : `str` or `NoneType`
        If being moved relatively, then this is the name of the column it is
        being moved relative to. If not a relative move then this will be
        `NoneType`.

    Notes
    -----
    The column addition may or may not have a default value and the placement
    may or may not be relative to an existing column.
    """
    try:
        # First is there an equals sign
        new_col, default_value = _process_equals(value)
    except ValueError as e:
        # No equals
        if 'unpack' not in e.args[0]:
            raise
        # TODO: return here
        return value, None, None, None

    # If we get here then we had an equals sign so now we need to see if we had
    # a relative column that we are adding next to
    try:
        # First is there an equals sign
        default_value, is_after, anchor_col = _process_gt_lt(default_value)
    except ValueError as e:
        # No equals
        if 'unpack' not in e.args[0]:
            raise
        # TODO: return here
        is_after = None
        anchor_col = None

    return new_col, default_value, is_after, anchor_col


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _process_equals(value):
    """Parse a transform that may contain an equals symbol indicating a column
    rename.

    Parameters
    ----------
    value : `str`
        A transform that potentially has a equals sign indicating a renaming of
        the column.

    Returns
    -------
    source_column : `str`
        The string that occured before the = symbol.
    dest_column : `str`
        The string that occured after the = symbol.
    """
    source_col, dest_col = [s.strip() for s in value.split('=')]
    return source_col, dest_col


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _process_gt_lt(value):
    """Parse a transform that may contain either a < or > symbol indicating a
    column relative move.

    Parameters
    ----------
    value : `str`
        A transform that potentially has a relative column movement. That is a
        < or > symbol separating a source column and an anchor column.

    Returns
    -------
    source_column : `str`
        The string that occured before the < or > symbol.
    is_after : `bool` or `NoneType`
        If the ``source_column`` is being moved after the anchor_col (`True`)
        or before it (`False`).
    anchor_col : `str` or `NoneType`
        This is the name of the column that the source column is being moved
        relative to.
    """
    source_column, direction, anchor_col = \
        [s.strip() for s in re.split(r'([<>])', value)]

    is_after = False
    if direction == '>':
        is_after = True
    return source_column, is_after, anchor_col


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