"""
Compress and index a file on a single column
"""
from pyaddons import __version__, __name__ as pkg_name, file_helper, gzopen
from Bio import bgzf
from simple_progress import progress
from contextlib import contextmanager
import pandas as pd
import tempfile
import builtins
import argparse
import sys
import os
import csv
import warnings
import shutil
import hashlib
import sqlite3
import pprint as pp


INDEX_SECTION = 'INDEX_SEEK'
METADATA_SECTION = 'METADATA'
MSG_PREFIX = '[info]'
MSG_OUT = sys.stderr


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main_fetch_index():
    """
    The main entry point to build an indexed wide file
    """
    # Initialise and parse the command line arguments
    parser = init_fetch_cmd_args()
    args = parse_fetch_cmd_args(parser)

    # Start a msg outputter, this will respect the verbosity and will output to
    # STDERR
    m = progress.Msg(prefix='', verbose=args.verbose, file=MSG_OUT)
    m.msg("= fetch from index ({0} v{1}) =".format(
        pkg_name, __version__))
    m.prefix = MSG_PREFIX
    m.msg_atts(args)

    csv_kwargs = {}
    if args.delimiter is not None:
        csv_kwargs['delimiter'] = args.delimiter
    else:
        csv_kwargs = get_csv_kwargs(
            args.infile,
            sniff_bytes=args.sniff_bytes,
            encoding=args.encoding
        )

    missing_values = []
    with open(args.infile, index_file=args.index_file,
              encoding=args.encoding, **csv_kwargs) as index:
        writer = csv.writer(sys.stdout, **csv_kwargs)

        for row in index.header:
            writer.writerow(row)

        for value in args.values:
            try:
                for row in index.fetch_value(value):
                    writer.writerow(row)
            except KeyError:
                missing_values.append(value)

    if len(missing_values) > 0:
        warnings.warn(
            "following values not in index: {0}".format(
                ",".join(missing_values)
                )
            )
    m.msg("*** END ***")


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def read_index(*args, dtypes={}, none_types=[''], **kwargs):
    """
    read from an index file into a pandas dataframe
    """
    infile = args[0]
    index_values = args[1:]

    if len(index_values) == 0:
        raise IndexError("no values to query")

    conn, infile_obj, reader, header_rows, index_metadata = \
        IndexReader.open_index(infile, **kwargs)
    infile_obj.close()

    try:
        if len(dtypes) == 0:
            idx = IndexReader(infile, none_types=none_types, **kwargs)
            idx.open()

            df = pd.DataFrame(
                _read_index_values(idx, index_values),
                columns=header_rows[-1]
            )
        else:
            class Dtypes(TypedIndexReader):
                EXPECTED_HEADER = header_rows[-1]
                OPT_CASTS = dtypes

            idx = Dtypes(infile, **kwargs)
            idx.open()

            df = pd.DataFrame(
                _read_index_values(idx, index_values)
            )
    finally:
        idx.close()
    return df


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _read_index_values(index_obj, values):
    missing_values = []
    for value in values:
        try:
            for row in index_obj.fetch_value(value):
                yield row
        except KeyError:
            missing_values.append(value)

    if len(missing_values) > 0:
        warnings.warn(
            "following values not in index: {0}".format(
                ",".join(missing_values)
                )
            )


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def init_fetch_cmd_args():
    """
    Initialise the command line arguments and return the parsed arguments

    Returns
    -------
    parser : :obj:`argparse.ArgumentParser`
        The argument parser with the arguments set up
    """
    parser = argparse.ArgumentParser(
        description="fetch rows matching index values from an indexed file"
    )

    # The wide format UKBB data fields filex
    parser.add_argument('infile', type=str, default='-',
                        help="The path to the indexed file")
    parser.add_argument('values', type=str, nargs='+',
                        help="index values to extract")
    parser.add_argument('-i', '--index-file', type=str,
                        help="An alternative path to the index file name. By "
                        "default the index file will be the same name as the "
                        "outfile but with an additional .idx extension")
    parser.add_argument('-v', '--verbose',  action="store_true",
                        help="Give more output")
    parser.add_argument('-d', '--delimiter',  type=str,
                        help="the file delimiter (default=detect)")
    parser.add_argument('-e', '--encoding',  type=str, default='utf-8',
                        help="the file encoding (default=utf-8)")
    parser.add_argument('-b', '--sniff-bytes',  type=int, default=2000000,
                        help="number of bytes at the start of the file to"
                        " sniff to detect the dialect of the file. If auto "
                        "detection fails then increase this value, if it "
                        "still fails use the API to index the file "
                        "(default=2000000 - (2MB))")

    return parser


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

    Parameters
    ----------
    parser : :obj:`argparse.ArgumentParser`
        The argument parser with the arguments set up

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

    # If no index file has been defined then we define it based on the
    # output file
    args.index_file = args.index_file or get_index_file_path(args.infile)

    # Make sure ~/ etc... is expanded
    for i in ['infile', 'index_file']:
        # Expand the file paths
        setattr(args, i, os.path.expanduser(getattr(args, i)))

    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_index_file_path(infile):
    """
    Get the potential index file path based on the input file name
    """
    return '{0}.idx'.format(infile)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class IndexReader(object):
    """
    A class for handling the read interaction with an indexed fat file
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @staticmethod
    def open_index(infile, index_file=None, sniff_bytes=2000000, cache=False,
                   encoding='utf-8', **csv_kwargs):
        """
        Open and validate an index file

        Parameters
        ----------
        infile : str
            The path to the indexed file
        index_file : str, optional, default: NoneType
            An alternative index file location. If `NoneType` then the index
            file path is the same as the wide file path but with an `.idx`
            extension
        sniff_bytes : int, optional, default: 2000000
            If no CSV kwargs have been defined, this is the number of bytes
            that are read in from the start of the file to attempt to detect
            the dialect
        encoding : str, optional, default: `utf-8`
            The encoding to open the indexed file with
        **csv_kwargs
            Arguments to csv for specifying delimiters etc...
        """
        # Get the correct CSV kwargs to open the file with
        csv_kwargs = get_csv_kwargs(
            infile,
            sniff_bytes=sniff_bytes,
            encoding=encoding,
            **csv_kwargs
        )

        # print(csv_kwargs['dialect'].delimiter == " ", "X")

        # open the indexed file
        infile_obj = bgzf.BgzfReader(infile)

        try:
            # Open the connection to the index database
            index_file = index_file or get_index_file_path(infile)
            conn = sqlite3.connect(index_file)
        except Exception:
            infile_obj.close()
            raise

        try:
            # Get all the metadata from the index so we can validate
            index_metadata = IndexReader.get_index_metadata(conn)

            # Now get the file MD5
            reader = csv.reader(infile_obj, **csv_kwargs)
            proc_header = ProcessHeader(
                reader,
                index_metadata['indexed_file_header_len']
            )

            # Cache the header rows as we will return them
            header_rows = [row for row in proc_header]
            if proc_header.md5 != index_metadata['header_md5']:
                raise KeyError(
                    "index header md5 and file md5 do not match: {0} != "
                    "{1}".format(proc_header.md5, index_metadata['header_md5'])
                )
        except Exception:
            infile_obj.close()
            conn.close()
            raise

        return conn, infile_obj, reader, header_rows, index_metadata

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @staticmethod
    def get_index_metadata(index):
        """
        Get the header MD5 as defined in the index

        Parameters
        ----------
        index : `sqlite3.Connection`
            The index connection

        Returns
        -------
        metadata : dict
            The metadata associated with the indexed file
        """
        curr = index.cursor()
        curr.execute(
            """
            SELECT header_md5,
                   indexed_file_header_len,
                   indexed_file_rows,
                   index_column,
                   expected_columns
            FROM index_metadata
            WHERE index_metadata_id = 1
            """
        )

        metadata = {}
        row = curr.fetchone()
        if row is None:
            raise IndexError("index metadata not defined")

        metadata['header_md5'] = row[0]
        metadata['indexed_file_header_len'] = row[1]
        metadata['indexed_file_rows'] = row[2]
        metadata['index_column'] = row[3]
        metadata['expected_columns'] = row[4]

        return metadata

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, infile, index_file=None, sniff_bytes=2000000,
                 encoding='utf-8', none_types=[], **csv_kwargs):
        """
        Parameters
        ----------
        infile : str
            The path to the indexed file
        index_file : str, optional, default: NoneType
            An alternative index file location. If `NoneType` then the index
            file path is the same as the wide file path but with an `.idx`
            extension
        sniff_bytes : int, optional, default: 2000000
            If no CSV kwargs have been defined, this is the number of bytes
            that are read in from the start of the file to attempt to detect
            the dialect
        cache : bool, optional, default: False
            By default, every time `fetch_value` is called then the index is
            queried to get the seek position. With this option set to `True`
            then the whole index is cached in memory first. This is designed to
            avoid lots of database queries at the expense of increased memory
            usage
        encoding : str, optional, default: `utf-8`
            The encoding to open the indexed file with
        **csv_kwargs
            Arguments to csv for specifying delimiters etc...
        """
        self._infile = os.path.expanduser(infile)
        self._index_file = index_file or get_index_file_path(infile)
        self._index_file = os.path.expanduser(self._index_file)
        self._sniff_bytes = sniff_bytes
        self._encoding = encoding
        self._csv_kwargs = csv_kwargs
        self._none_types = none_types

        # As we have overridden open we have to use the builtin
        builtins.open(self._infile).close()
        builtins.open(self._index_file).close()

        self._opened = False
        self._conn = None
        self._cursor = None
        self._infile_obj = None
        self._reader = None
        self._header_rows = []
        self._cache = {}

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """
        Open the wide file and the index
        """
        self._cache = {}
        self._check_closed()
        self._conn, self._infile_obj, self._reader, self._header_rows, \
            self._index_metadata = IndexReader.open_index(
                self._infile,
                index_file=self._index_file,
                sniff_bytes=self._sniff_bytes,
                encoding=self._encoding,
                **self._csv_kwargs
            )

        self._cursor = self._conn.cursor()

        if self._cache is True:
            self._cache_index()

        self._opened = True

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """
        Close the config file and the wide file
        """
        try:
            self._infile_obj.close()
        except AttributeError:
            pass

        try:
            self._conn.close()
        except AttributeError:
            pass

        self._idx = None
        self._opened = False

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def seek(self, to):
        """
        Expose the underlying file objects seek method
        """
        self._infile_obj.seek(to)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def values(self):
        """
        Return the a list of values in the index
        """
        self._cursor.execute(
            """
            SELECT value FROM index_data
            """
        )

        return [i[0] for i in self._cursor.fetchall()]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def n_values(self):
        """
        Return the number of samples in the index
        """
        self._cursor.execute(
            """
            SELECT count(*) FROM index_data
            """
        )

        return self._cursor.fetchone()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def header(self):
        """
        Return the a list of values in the index
        """
        return self._header_rows

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _cache_index(self):
        """
        Cache the whole contents of the index into memory
        """
        self._cursor.execute(
            """
            SELECT value, seek_pos, nlines
            FROM index_data
            """
        )
        for row in self._cursor:
            self._cache[row[0]] = row

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _check_open(self):
        """
        Error out if the files are not open
        """
        if self._opened is False:
            raise IOError("index not open")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _check_closed(self):
        """
        Error out if the files are open
        """
        if self._opened is True:
            raise IOError("index is open")

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def fetch_value(self, value):
        """
        Fetch an un-formatted row for a single value

        Parameters
        ----------
        eid : int
            A single UKBB sample ID

        Returns
        -------
        row : list of str
            A single un-formatted row

        Raises
        ------
        KeyError
            If the sample EID is not in the index
        RuntimeError
            If the first column of the extracted row is not the same as the
            value that has been requested
        """
        value = str(value)
        try:
            index_row = self._cache[value]
        except KeyError:
            try:
                self._cursor.execute(
                    """
                    SELECT value, seek_pos, nlines
                    FROM index_data
                    WHERE value = ?
                    """,
                    (value, )
                )
                index_row = self._cursor.fetchone()
            except AttributeError as e:
                raise AttributeError("have you opened the index?") from e

        if index_row is None:
            raise KeyError("'{0}' not found in index".format(value))

        for row in self._seek_row(value, index_row):
            yield [i if i not in self._none_types else None for i in row]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _seek_row(self, value, index_row):
        """
        Fetch an un-formatted row for a single value

        Parameters
        ----------
        eid : int
            A single UKBB sample ID

        Returns
        -------
        row : list of str
            A single un-formatted row

        Raises
        ------
        KeyError
            If the sample EID is not in the index
        RuntimeError
            If the first column of the extracted row is not the same as the
            value that has been requested
        """
        self._infile_obj.seek(index_row[1])

        for i in range(index_row[2]):
            row = next(self._reader)
            if row[self._index_metadata['index_column']] != value:
                raise RuntimeError(
                    "bad index: expected {0}, found '{1}'".format(
                        value, row[self._index_metadata['index_column']]
                    )
                )
            if len(row) != self._index_metadata['expected_columns']:
                raise IndexError(
                    "{0}: row ('{1}') is not the same length as"
                    " header ('{2}') ".format(
                        value,
                        len(row),
                        self._index_metadata['expected_columns']
                    )
                )
            yield row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
class TypedIndexReader(IndexReader):
    """
    An index reader that will return fetched rows as dictionaries where the
    keys are column names and the values are formatted to the correct type
    base defined by the user. This should be sub-classed and the subclass
    should define column name -> datatype mappings
    """
    EXPECTED_HEADER = []

    # Keys should be column names and the values should be cast functions. If
    # these casts fail then an error will be raised
    CASTS = {}

    # Keys should be column names and the values should be optional cast
    # functions. If these casts fail then an error will not be raised
    OPT_CASTS = {}

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        """
        """
        super().__init__(*args, **kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def header(self):
        """
        Return the header row, this is the last row in all the header rows
        """
        return self._header_rows[-1]

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """
        Open the wide file and the index
        """
        # Open with the super class and then make sure that we have some
        # header lines. The actual header is treated as the last header_row
        super().open()

        # Make sure we have enough data in the header to do something
        if self._header_rows == 0:
            raise ValueError(
                "'{0}' requires at least one header row".format(
                    self.__class__.__name__)
            )

        # Make sure that the expected header matches what we actually have
        for i in range(len(self.__class__.EXPECTED_HEADER)):
            if self.__class__.EXPECTED_HEADER[i] != self.header[i]:
                raise KeyError(
                    "expected value '{0}' does not match actual value '{1}' "
                    "in the header".format(
                        self.__class__.EXPECTED_HEADER[i],
                        self.header[i]
                    )
                )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def fetch_value(self, value):
        """
        Fetch an un-formatted row for a single value

        Parameters
        ----------
        eid : int
            A single UKBB sample ID

        Returns
        -------
        row : list of str
            A single un-formatted row

        Raises
        ------
        KeyError
            If the sample EID is not in the index
        RuntimeError
            If the first column of the extracted row is not the same as the
            value that has been requested
        """
        # call the super class and then format the row
        for row in super().fetch_value(value):
            yield self._format_row(
                self._row_as_dict(
                    make_none_type(row)
                )
            )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _row_as_dict(self, row):
        """
        Make the row into the dictionary

        Parameters
        ----------
        row : list of str
            The data row as read in by `csv.Reader`

        Returns
        -------
        row : dict
            The row formatted as a dict, each key is a column name and each
            value is a data value
        """
        return dict([(col_name, row[idx])
                     for idx, col_name in enumerate(self.header)])

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def _format_row(self, row):
        """
        Format the row. This checks the row length is what is required and then
        will perform any casts or specific data reformating on the row

        Parameters
        ----------
        row : list of str
            The data row to process

        Raises
        ------
        IndexError
            If the length of the data row is not what is expected
        """
        for key, cast_func in self.__class__.CASTS.items():
            row[key] = cast_func(row[key])

        for key, cast_func in self.__class__.OPT_CASTS.items():
            try:
                row[key] = cast_func(row[key])
            except Exception as e:
                # print(e)
                pass

        return row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def make_none_type(row):
    """
    For a single row make '' into `NoneType`

    Parameters
    ----------
    row : list or dict
        A single data row with missing data represented as ''

    Returns
    -------
    row : list
        A single data row with '' made into `NoneType`

    """
    return [None if i.strip() == '' else i for i in row]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@contextmanager
def open(*args, **kwargs):
    """
    A context wrapper for the index reading

    Parameters
    ----------
    *args
        Arguments supplied to `wide_file_index.IndexReader`
    **kwargs
        Keyword arguments supplied to `wide_file_index.IndexReader`
    """

    # Code to acquire resource, e.g.:
    idx = IndexReader(*args, **kwargs)
    try:
        # Open the index and supply it
        idx.open()
        yield idx
    finally:
        # close the index
        idx.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main_create_index():
    """
    The main entry point to build an indexed wide file
    """
    # Initialise and parse the command line arguments
    parser = init_create_cmd_args()
    args = parse_create_cmd_args(parser)

    # Start a msg outputter, this will respect the verbosity and will output to
    # STDERR
    m = progress.Msg(prefix='', verbose=args.verbose, file=MSG_OUT)
    m.msg("= file indexer ({0} v{1}) =".format(
        pkg_name, __version__))
    m.prefix = MSG_PREFIX
    m.msg_atts(args)

    csv_kwargs = {}
    # Make sure the wide file can be opened only if input is not from STDIN
    if args.file_path != '-':
        builtins.open(args.file_path).close()
    elif args.delimiter is None:
        # if input is from STDIN then we have to set the delimiter if not
        # already set as we can't sniff the csv args
        args.delimiter = "\t"

    # if the delimiter has been set then update the csv keyword arguments
    if args.delimiter is not None:
        csv_kwargs['delimiter'] = args.delimiter

    try:
        builtins.open(args.outfile).close()
        raise FileExistsError("index already exists exists, use --clobber")
    except FileNotFoundError:
        pass
    except FileExistsError:
        if args.clobber is True:
            os.unlink(args.outfile)
        else:
            raise

    # Initialise the building of the index file. When used on the command line
    # the default is to sniff the csv dialect
    build_index_file(
        args.file_path,
        args.outfile,
        args.index_file,
        skip_lines=args.skip_lines,
        index_col=args.index_col,
        sniff_bytes=args.sniff_bytes,
        verbose=args.verbose,
        key_size=args.key_size,
        dir=args.tmp_dir,
        encoding=args.encoding,
        **csv_kwargs
    )

    m.msg("*** END ***")


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def init_create_cmd_args():
    """
    Initialise the command line arguments and return the parsed arguments

    Returns
    -------
    parser : :obj:`argparse.ArgumentParser`
        The argument parser with the arguments set up
    """
    parser = argparse.ArgumentParser(
        description="build a compressed indexed file for random access. The "
        " file must be presorted on the index column"
    )

    # The wide format UKBB data fields filex
    parser.add_argument('file_path', type=str, default='-',
                        help="The path to the file file, input from STDIN "
                        "use - for input file")
    parser.add_argument('outfile', type=str,
                        help="An output indexed file path. The output file "
                        "will be bgzipped, as part of the indexing process")
    parser.add_argument('-i', '--index-file', type=str,
                        help="An alternative path to the index file name. By "
                        "default the index file will be the same name as the "
                        "outfile but with an additional .idx extension")
    parser.add_argument('-c', '--index-col', type=int, default=1,
                        help="The column number (1-based) that should be "
                        "indexed. Note that the file should be pre-sorted on "
                        "this column as well (default=1)")
    parser.add_argument('-v', '--verbose',  action="store_true",
                        help="Give more output")
    parser.add_argument('-l', '--clobber',  action="store_true",
                        help="Delete any existing indexes")
    parser.add_argument('-s', '--skip-lines',  type=int, default=1,
                        help="number of header lines to skip (default=1)")
    parser.add_argument('-k', '--key-size',  type=int, default=2000,
                        help="The max number of characters a key entry can be."
                        " Index values above this length will generate errors,"
                        " however, larger key-size will generate bigger "
                        "index files (default=2000)")
    parser.add_argument('-d', '--delimiter',  type=str,
                        help="the file delimiter (default=TAB)")
    parser.add_argument('-e', '--encoding',  type=str, default='utf-8',
                        help="the file encoding (default=utf-8)")
    parser.add_argument('-t', '--tmp-dir',  type=str,
                        help="an alternative tmp-dir location "
                        "(default=system tmp)")
    parser.add_argument('-b', '--sniff-bytes',  type=int, default=2000000,
                        help="number of bytes at the start of the file to"
                        " sniff to detect the dialect of the file. If auto "
                        "detection fails then increase this value, if it "
                        "still fails use the API to index the file "
                        "(default=2000000 - (2MB))")

    return parser


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

    Parameters
    ----------
    parser : :obj:`argparse.ArgumentParser`
        The argument parser with the arguments set up

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

    # If no index file has been defined then we define it based on the
    # output file
    if args.index_file is None:
        args.index_file = '{0}.idx'.format(args.outfile)

    # Make sure ~/ etc... is expanded
    for i in ['file_path', 'index_file', 'outfile']:
        # If input is from STDIN then do not check
        if i == 'file_path' and getattr(args, i) == '-':
            continue

        # Expand the file paths
        setattr(args, i, os.path.expanduser(getattr(args, i)))

    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def build_index_file(file_path, outfile, index_file, skip_lines=1, index_col=1,
                     sniff_bytes=2000000, verbose=False, check_sort=True,
                     commit_every=100000, key_size=2000, encoding='utf-8',
                     **kwargs):
    """
    Build a compressed indexed file and associated index file (which is
    actually an SQLite database)

    Parameters
    ----------
    file_path : str
        The location of the file to be indexed, this may or may not be
        compressed. If it is a dash (-) then input from STDIn is expected
    outfile : str
        The location of the BGZIP compressed output file that will have been
        indexed. The outfile is allowed to be the same as the input wide file
    index_file : str
        The location of the index file. This should not be the same as any of
        the output or wide file paths
    skip_lines : int, optional, default: 1
        The number of header lines to skip (not index) at the beginning of the
        file
    index_col : int, optional, default: 1
        A 1-based column number that will be indexed. It is assumed that the
        input file (`file_path`), is pre-sorted on this column and an error
        will be raised if not (or more precisely if an index value is seen >
        1 time in different areas of the file)
    sniff_bytes : int, optional, default: 2000000
        If no csv_kwargs have been supplied then we attempt to sniff out the
        delimiter of the file. This is the number of bytes that is used to
        determine the dialect. Given the file is wide, a value of 2MB has been
        chosen but can be increased if not sufficient
    check_sort : bool, optional, default: True
        Ensure the file is sorted properly, this will cache the index values
        in a set and may take excess memory if there are loads of values. Set
        this to False to disable those checks
    verbose : bool, optional, default: False
    **kwargs
        Any arguments that should be used by csv to read the wide file being
        indexed. These are also recycled into the arguments to write the
        indexed file. If not supplied then the dialect is sniffed out. In
        addition, any arguments accepted by `tempfile.mkstemp`
    """
    index_col = index_col - 1

    # Extract the arguments for tempfile from all the kwargs
    tempfile_kwargs = _extract_tempfile_args(kwargs)

    # Now get the csv keyword arguments only if input is not from STDIN
    if file_path != '-':
        kwargs = get_csv_kwargs(file_path, encoding=encoding,
                                sniff_bytes=sniff_bytes, **kwargs)

    # Get a temp index file
    temp_index_conn, temp_index_name = get_temp_index(
        key_size=key_size,
        **tempfile_kwargs
    )
    temp_index_curr = temp_index_conn.cursor()

    try:
        # Block compression output file
        temp_out_name = _get_temp_file(**tempfile_kwargs)
        # write_obj = open(temp_out_name, 'wt', encoding=encoding)
        outbgz = bgzf.BgzfWriter(temp_out_name)
        # outbgz = bgzf.BgzfWriter(fileobj=write_obj)

        # Now make sure we have set up the input file correctly, the user must
        # supply a - if they are giving input from STDIN
        if file_path == '-':
            sys.stdin.reconfigure(encoding=encoding)
            infile = sys.stdin
        else:
            # GZIP agnostic, so does not care if compressed or uncompressed
            infile = gzopen.gzip_agnostic(file_path, 't', encoding=encoding)

        reader = csv.reader(infile, **kwargs)
        writer = csv.writer(outbgz, lineterminator=os.linesep, **kwargs)

        prog = progress.RateProgress(verbose=verbose, file=MSG_OUT)

        proc_header = ProcessHeader(reader, skip_lines)
        for header_row in proc_header:
            writer.writerow(header_row)
        header_md5 = proc_header.md5

        seen_idx = set()
        curr_idx = -1
        curr_idx_row = None
        nlines = 0
        batch = []
        idx_id = 1
        cur_ncols = -1
        for row_no, row in enumerate(prog.progress(reader), 1):
            if cur_ncols != len(row):
                if row_no > 1:
                    raise IndexError(
                        "number of columns differ at line {0}: {1} cols vs {2}"
                        " cols".format(row_no, cur_ncols, len(row))
                    )
            cur_ncols = len(row)
            if curr_idx != row[index_col]:
                try:
                    curr_idx_row.append(nlines)
                    batch.append(tuple(curr_idx_row))
                except AttributeError:
                    # First row
                    pass

                curr_idx = row[index_col]

                if check_sort is True and curr_idx in seen_idx:
                    raise KeyError(
                        "Already seen index entry, is your file sorted?:"
                        " {0}".format(curr_idx)
                    )

                if len(curr_idx) > key_size:
                    raise KeyError(
                        "Index entry is too long - please increase index size"
                    )

                seen_idx.add(curr_idx)
                nlines = 0
                curr_idx_row = [idx_id, curr_idx, outbgz.tell()]
                idx_id += 1

            if (commit_every % row_no) == 0:
                _insert_index_batch(temp_index_conn, temp_index_curr, batch)
                batch = []

            try:
                writer.writerow(row)
            except UnicodeEncodeError:
                for idx, i in enumerate(row):
                    try:
                        row[idx] = i.encode().decode(
                            'latin1', 'replace'
                        )
                    except UnicodeEncodeError:
                        row[idx] = i.encode().decode(
                            'latin1', 'backslashreplace'
                        )
                        warnings.warn(
                            "backslash replace on row: '{0}'".format(row_no)
                        )
                writer.writerow(row)
            nlines += 1

        # Append the data for the last index in the file
        curr_idx_row.append(nlines)
        batch.append(tuple(curr_idx_row))
        _insert_index_batch(temp_index_conn, temp_index_curr, batch)
        batch = []

        # Now append the
        temp_index_curr.execute(
            '''
            INSERT INTO index_metadata(
                index_metadata_id,
                header_md5,
                index_column,
                expected_columns,
                indexed_file_header_len,
                indexed_file_rows
            )
            VALUES (?, ?, ?, ?, ?, ?)
            ''', (1, header_md5, index_col, cur_ncols, skip_lines, row_no)
        )
        temp_index_conn.commit()
    except Exception:
        temp_index_conn.close()
        outbgz.close()
        infile.close()
        os.unlink(temp_index_name)
        os.unlink(temp_out_name)
        raise

    # Close databases and files
    temp_index_conn.close()
    outbgz.close()
    infile.close()

    # Relocate to the final non-temp location
    shutil.move(temp_index_name, index_file)
    shutil.move(temp_out_name, outfile)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_tempfile_args(kwargs):
    """
    Extract the tempfile arguments from an dict of all kwargs arguments

    Parameters
    ----------
    kwargs : dict
        A dictionary of keyword arguments to extract the tempfile arguments
        from

    Returns
    -------
    tempfile_kwargs : dict
        A dictionary of tempfile keyword arguments. If there are none then this
        will be an empty dictionary
    """
    tempfile_kwargs = {}
    for arg in ['dir', 'suffix', 'prefix', 'text']:
        try:
            tempfile_kwargs[arg] = kwargs.pop(arg)
        except KeyError:
            pass

    return tempfile_kwargs


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_csv_kwargs(file_path, sniff_bytes=2000000, encoding='utf-8', **kwargs):
    """
    Sniff out csv dialect if needed

    Parameters
    ----------
    kwargs : dict
        Existing csv keywords argument, if this has some length then no
        sniffing is needed

    Returns
    -------
    kwargs : dict
        The keyword arguments with the dialect set if needed. Note that the
        return is not strictly necessary as the dialect is added in place
    """
    # Make sure that we sniff out the csv args if they have not been supplied
    if len(kwargs) == 0:
        kwargs['dialect'] = file_helper.get_dialect(
            file_path,
            sniff_bytes=sniff_bytes,
            encoding=encoding,
            **kwargs
        )
    return kwargs


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_temp_index(key_size=2000, **kwargs):
    """
    This generates a temp index file (which is actually an SQLite database)

    Parameters
    ----------
    **kwargs
        Arguments to `tempfile.mkstemp`

    Returns
    -------
    index_file : `sqlalchemy.Sessionmaker`
        A connection to the database that makes up the index file
    index_file_name : str
        The full path to the temp index file
    """
    temp_file_name = _get_temp_file(**kwargs)
    conn = sqlite3.connect(temp_file_name)
    _create_tables(conn, key_size=key_size)
    return conn, temp_file_name


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _create_tables(conn, key_size=2000):
    """

    """
    try:
        key_size = int(key_size)
        if key_size <= 0:
            raise TypeError("too small")
    except TypeError as e:
        raise ValueError("invalid key size: {0}".format(key_size)) from e

    conn.execute("""
    CREATE TABLE index_metadata
    (
    index_metadata_id INTEGER PRIMARY KEY,
    header_md5 VARCHAR(32) NOT NULL,
    index_column INTEGER NOT NULL,
    expected_columns INTEGER NOT NULL,
    indexed_file_header_len INTEGER NOT NULL,
    indexed_file_rows INTEGER NOT NULL
    ) WITHOUT ROWID;
    """)

    # I am not sure that I can bind the key size properly so I am using
    # format
    conn.execute("""
    CREATE TABLE index_data
    (
    index_data_id INTEGER PRIMARY KEY,
    value VARCHAR({0}) NOT NULL,
    seek_pos INTEGER NOT NULL,
    nlines INTEGER NOT NULL
    ) WITHOUT ROWID;
    """.format(key_size))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _get_temp_file(**kwargs):
    """
    Get the path to a temp file

    Parameters
    ----------
    **kwargs
        Arguments to `tempfile.mkstemp`

    Returns
    -------
    temp_file_name : str
        The path to a temp file
    """
    temp_file_obj, temp_file_name = tempfile.mkstemp(**kwargs)
    os.close(temp_file_obj)
    return temp_file_name


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
class ProcessHeader(object):
    def __init__(self, reader, nlines):
        self._reader = reader
        self._nlines = nlines
        self._hash_md5 = hashlib.md5()
        self._line_count = 0

    @property
    def md5(self):
        """

        """
        if self._nlines == 0:
            self._hash_md5.update(b'')
        return self._hash_md5.hexdigest()

    def __iter__(self):
        return self

    def __next__(self):
        if self._line_count == self._nlines:
            raise StopIteration("finished")

        try:
            line = next(self._reader)
        except StopIteration:
            raise IOError("no data in file after processing header")

        self._hash_md5.update("|".join(line).encode())
        self._line_count += 1
        return line


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _insert_index_batch(conn, cur, batch):
    """
    """
    # print("CALLED")
    if len(batch) > 0:
        cur.executemany(
            '''
            INSERT INTO index_data(index_data_id, value, seek_pos, nlines)
            VALUES (?,?,?,?)
            ''', batch
        )
    conn.commit()


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