"""Classes for writing content to sorted intermediate files

There are a several classes providing the same interface just different
implementations so I can understand what technique is best.
"""
from merge_sort import common
from pyaddons import utils
from sortedcontainers import SortedList
import gzip
import os
import csv
import pprint as pp


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class _BaseChunks(object):
    """Base class for Write files that are sorted on key where each file has
    at max chunksize rows.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, chunk_dir, key=None, chunksize=100000, write_mode='wt',
                 write_method=gzip.open, header=None, delete_on_error=True,
                 reverse=False, chunk_prefix=None, chunk_suffix=".chunk.gz",
                 **write_kwargs):
        self._chunk_dir = chunk_dir
        self._key = key
        self._chunksize = None
        self._write_mode = write_mode
        self._write_method = write_method
        self._header = header
        self._delete_on_error = delete_on_error
        self._reverse = None
        self._chunk_prefix = chunk_prefix
        self._chunk_suffix = chunk_suffix
        self._write_kwargs = write_kwargs
        self.reverse = reverse
        self.chunksize = chunksize

        self._chunks = {}
        self._tmp_file = None
        self._cur_chunk = self.get_empty_chunk()
        self._idx = 0
        self._n_added = 0
        self._n_written = 0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __enter__(self):
        """Entry point for the context manager.
        """
        return self.open()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __exit__(self, *args):
        """Exit point for the context manager

        Parameters
        ----------
        *args
            A tuple containing any errors that may have caused the context
            manager to exit early. If no errors were generated and the context
            manager exited normally the tuple will be (None, None, None).
        """
        self.close()

        if args[0] is not None and self._delete_on_error is True:
            # TODO: Delete all the chunk files if we error
            self.delete()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open(self):
        """Open the input file and setup everything for reading.

        Returns
        -------
        self : `merge_sort.chunks._BaseChunks`
            For chaining.
        """
        self.open_chunk_file()
        self.write_header()
        return self

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close all files.
        """
        self.flush_chunk()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def delete(self):
        """Delete all the chunk files created by the object.

        If the chunk file does not exist, this will fail silently, other errors
        will be raised though.
        """
        for i in self.chunk_file_list:
            try:
                os.unlink(i)
            except FileNotFoundError:
                # Might have already been deleted by something else, if used in
                # an external merge sort
                pass

        self._chunks = {}

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_row(self, row):
        """Add a row to the current chunk.

        Parameters
        ----------
        row : `str`
            The row to add to the file. Should contain a newline at the end.
        """
        self._n_added += 1

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunk_files(self):
        """Get all the chunk file names (`dict`).
        """
        return self._chunks

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunk_file_list(self):
        """Get all the chunk file names in the order they were written
        (`list` of `str`).
        """
        return list(self._chunks.keys())

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def n_chunks(self):
        """Get the number of chunks that have been written to file (this does
        not include the current chunk) (`int`)
        """
        return len(self._chunks)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def n_added(self):
        """Get the number of rows in the current chunk (`int`).
        """
        return self._n_added

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def n_written(self):
        """Get the number of rows in the current chunk (`int`).
        """
        return self._n_written

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def n_written_total(self):
        """Get the total number of rows written to all chunks
        (not including any current chunks) (`int`).
        """
        return sum(self._chunks.values())

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def current_chunk_file(self):
        """Get the current chunk file object (`File`).
        """
        return self._tmp_file

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def current_chunk(self):
        """Get the current chunk (`list`).
        """
        return self._cur_chunk

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def reverse(self):
        """Get the reverse argument (`bool`).
        """
        return self._reverse

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def header(self):
        """Get the header argument (`bool`).
        """
        return self._header

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @reverse.setter
    def reverse(self, reverse):
        """Set the reverse argument (`bool`).
        """
        if not isinstance(reverse, bool):
            raise TypeError(
                f"reverse should be a Boolean, not: {type(reverse)}"
            )
        self._reverse = reverse

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def chunksize(self):
        """Get the chunksize argument, should be a positive integer (`int`).
        """
        return self._chunksize

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @chunksize.setter
    def chunksize(self, chunksize):
        """Set the chunksize argument, should be a positive integer (`int`).
        """
        if not isinstance(chunksize, int):
            raise TypeError(
                f"chunksize should be a positive integer, not: {chunksize}"
            )

        if chunksize <= 0:
            raise ValueError(
                f"chunksize should be a positive integer, not: {chunksize}"
            )
        self._chunksize = chunksize

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @property
    def key(self):
        """Get the sort key argument (`function`).
        """
        return self._key

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @chunksize.setter
    def chunksize(self, chunksize):
        """Set the chunksize argument, should be a positive integer (`int`).
        """
        if not isinstance(chunksize, int):
            raise TypeError(
                f"chunksize should be a positive integer, not: {chunksize}"
            )

        if chunksize <= 0:
            raise ValueError(
                f"chunksize should be a positive integer, not: {chunksize}"
            )
        self._chunksize = chunksize

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def set_new_chunk_file(self):
        """Initialise a new chunk file to write the chunk to.

        Raises
        ------
        FileExistsError
            If any existing chunk files have not been closed first.
        """
        if self._tmp_file is not None:
            raise FileExistsError("please close existing chunk file first")

        # Create a new one
        self._tmp_file = utils.get_tmp_file(
            dir=self._chunk_dir, prefix=self._chunk_prefix,
            suffix=self._chunk_suffix
        )

        self._tmp_file = self._write_method(
            self._tmp_file, self._write_mode, **self._write_kwargs
        )

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_chunk_file(self):
        """Open a chunk file.

        Returns
        -------
        chunk_file: `File`
            The file object for the chunk.

        Raises
        ------
        FileExistsError
            If any existing chunk files have not been closed first.
        """
        self.set_new_chunk_file()
        return self._tmp_file

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close_chunk_file(self):
        """Close a tmp chunk file

        If one is not initialise then it will fail silently. The closed chunk
        file name will be added to the _chunks dict only if it has been written
        to.
        """
        try:
            self._tmp_file.close()
        except AttributeError:
            return

        if self._n_written > 0:
            self._chunks[self._tmp_file.name] = self._n_written
        else:
            # remove the empty chunk file
            os.unlink(self._tmp_file.name)

        # Empty the file
        self._tmp_file = None
        self._n_written = 0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_header(self):
        """write the header row to the chunk file.
        """
        if self._header is not None:
            self.write_row(self.header)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def sort_chunk(self):
        """Sort the chunk.
        """
        pass

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_empty_chunk(self):
        """Get an empty chunk holder.

        Returns
        -------
        empty_chunk : `list`
            An empty chunk to fill.
        """
        return []

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def clear_chunk(self):
        """Clear the contents of the chunk.
        """
        self._cur_chunk = self.get_empty_chunk()
        self._n_added = 0

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_chunk(self):
        """Write the chunk to file.
        """
        for row in self._cur_chunk:
            self.write_row(row)
            self._n_written += 1

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_row(self, row):
        """Write a single row to the chunk file.

        Parameters
        ----------
        row : `str`
            The row to write to the file. The row should contain a newline
            character.
        """
        self._tmp_file.write((str(row)))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flush_chunk(self):
        """Perform the whole process of sorting,opening,wring,closing and
        clearing a chunk.
        """
        pass


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class SortedChunks(_BaseChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def sort_chunk(self):
        """Sort the chunk.
        """
        self.current_chunk.sort(key=self.key, reverse=self.reverse)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_row(self, row):
        """Add a row to the current chunk.

        Parameters
        ----------
        row : `str`
            The row to add to the file. Should contain a newline at the end.
        """
        super().add_row(row)
        self.current_chunk.append(row)

        if self.n_added == self.chunksize:
            self.flush_chunk()
            self.open_chunk_file()
            self.write_header()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flush_chunk(self):
        """Perform the whole process of sorting,opening,wring,closing and
        clearing a chunk.
        """
        self.sort_chunk()
        self.write_chunk()
        self.close_chunk_file()
        self.clear_chunk()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class SortedListChunks(_BaseChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows.

    This implementation uses sortedcontainers rather than lists.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    """

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_empty_chunk(self):
        """Get an empty chunk holder.

        Returns
        -------
        empty_chunk : `list`
            An empty chunk to fill.
        """
        return SortedList([], key=self.key)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_row(self, row):
        """Add a row to the current chunk.

        Parameters
        ----------
        row : `str`
            The row to add to the file. Should contain a newline at the end.
        """
        super().add_row(row)
        self.current_chunk.add(row)

        if self.n_added == self.chunksize:
            self.flush_chunk()
            self.open_chunk_file()
            self.write_header()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flush_chunk(self):
        """Perform the whole process of sorting,opening,wring,closing and
        clearing a chunk.
        """
        self.sort_chunk()
        self.write_chunk()
        self.close_chunk_file()
        self.clear_chunk()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class ExtendedChunks(_BaseChunks):
    """Write files that are sorted on key where each file can have variable
    number of rows possibly exceeding the chunksize value.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.

    Notes
    -----
    This will keep adding chunks to a chunk file if the sort order is
    maintained between the last element of the previous chunk and the first
    element of the current chunk. This can reduce the number of chunk files
    that are needed and is useful if you suspect that the input may be sorted
    already.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._last_write = None

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def close(self):
        """Close the chunker and any open chunk files.
        """
        super().close()
        self.close_chunk_file()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def sort_chunk(self):
        """Sort the chunk.
        """
        self._cur_chunk.sort(key=self.key, reverse=self.reverse)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_row(self, row):
        """Add a row to the current chunk.

        Parameters
        ----------
        row : `str`
            The row to add to the file. Should contain a newline at the end.
        """
        super().add_row(row)
        self.current_chunk.append(row)

        if self.n_added == self.chunksize:
            self.flush_chunk()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def flush_chunk(self):
        """Perform the whole process of sorting,opening,wring,closing and
        clearing a chunk.
        """
        self.sort_chunk()

        try:
            # Put together a list to evaluate if this chunk continues from the
            # previous one
            ref = [self._last_write, self._cur_chunk[0]]
        except IndexError:
            # The index error probably means that the current chunk is empty so
            # we check that and if not we will raise
            if len(self.current_chunk) > 0:
                raise
            # Nothing left to do as there is no data in the current chunk
            return

        try:
            result = sorted(ref, key=self.key, reverse=self.reverse)
        except Exception:
            # The type errpr is raised if the last write is not initialised
            # which will happen on the first write to the first file, but we
            # will check to make sure
            if self._last_write is not None:
                raise
            result = ref

        # Is the sorted list the same as the boundry list
        # TODO: Should I just compare the first element? Better to compare the
        #  actual keys
        if ref == result:
            # extend
            self.write_chunk()
        else:
            # write a new file
            self.close_chunk_file()
            self.open_chunk_file()
            self.write_header()
            self.write_chunk()

        self._last_write = self.current_chunk[-1]
        self.clear_chunk()


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvMixin(object):
    """Add CSV writing abilities to the various chunkers.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def initialise(self, kwargs):
        """Initialise the mixin with the CSV keyword arguments.
        """
        self.writer = None
        self.csv_kwargs = common.extract_csv_kwargs(kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_chunk_file(self):
        """Open a chunk file.

        Returns
        -------
        chunk_filename : `str`
            The name of the chunk file that has been opened.

        Raises
        ------
        FileExistsError
            If any existing chunk files have not been closed first.
        """
        self.set_new_chunk_file()
        self.writer = csv.writer(self.current_chunk_file, **self.csv_kwargs)
        return self.current_chunk_file

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_row(self, row):
        """Write a row to the current chunk.

        Parameters
        ----------
        row : `list`
            The row to add to the file.
        """
        self.writer.writerow(row)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvSortedChunks(CsvMixin, SortedChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows. This uses a csv.writer so rows should be added as lists.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvSortedListChunks(CsvMixin, SortedListChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows.

    This implementation uses sortedcontainers rather than lists. This uses a
    csv.writer so rows should be added as lists.


    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvExtendedChunks(CsvMixin, ExtendedChunks):
    """Write files that are sorted on key where each file can have variable
    number of rows possibly exceeding the chunksize value.

    This uses a csv.writer so rows should be added as lists.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An optional header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.

    Notes
    -----
    This will keep adding chunks to a chunk file if the sort order is
    maintained between the last element of the previous chunk and the first
    element of the current chunk. This can reduce the number of chunk files
    that are needed and is useful if you suspect that the input may be sorted
    already.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvDictMixin(object):
    """Add CSV dict writer abilities to the various chunkers.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def initialise(self, kwargs):
        """Initialise the mixin with the CSV dict writer keyword arguments.
        """
        self.writer = None
        self.csv_kwargs = common.extract_csv_kwargs(kwargs)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def open_chunk_file(self):
        """Open a chunk file.

        Returns
        -------
        chunk_filename : `str`
            The name of the chunk file that has been opened.

        Raises
        ------
        FileExistsError
            If any existing chunk files have not been closed first.
        """
        if self.header is None:
            raise ValueError(
                "The dictionary mixin requires a header to be available"
            )

        self.set_new_chunk_file()
        self.writer = csv.DictWriter(
            self.current_chunk_file, self.header, **self.csv_kwargs
        )
        return self.current_chunk_file

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_header(self):
        """Write a header to the file.
        """
        self.writer.writeheader()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def write_row(self, row):
        """Write a row to the current chunk.

        Parameters
        ----------
        row : `dict`
            The row to add to the file.
        """
        try:
            self.writer.writerow(row)
        except Exception:
            pp.pprint(row)
            raise


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvDictSortedChunks(CsvDictMixin, SortedChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows.

    This uses a ``csv.DictWriter`` so rows should be added as dictionaries. The
    ``header`` is absolutely required.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An required header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvDictSortedListChunks(CsvDictMixin, SortedListChunks):
    """Write files that are sorted on key where each file has at max
    chunksize rows.

    This implementation uses sortedcontainers rather than lists. This uses a
    ``csv.DictWriter`` so rows should be added as dictionaries. The ``header``
    is absolutely required.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        A required header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class CsvDictExtendedChunks(CsvDictMixin, ExtendedChunks):
    """Write files that are sorted on key where each file can have variable
    number of rows possibly exceeding the chunksize value.

    This uses a ``csv.DictWriter`` so rows should be added as dictionaries. The
    ``header`` is absolutely required.

    Parameters
    ----------
    chunk_dir : `str`
        The directory to write the chunk files to.
    key : `function`
        A key function that will return the values to sort on. This will be
        passed onto the sorting routine.
    chunksize : `int`, optional, default: `100000`
        The number of rows to store in memory before sorting and outputting to
        the chunk file.
    write_mode : `function`, optional, default: `wt`
        The method mode to use to write the chunk files. In most cases this
        should not be changed but it is there just in case the ``write_method``
        has any odd requirements.
    write_method : `function`, optional, default: `gzip.open`
        The method to use to write the chunk files.
    header : `list` of `str`, optional, default: `NoneType`
        An required header to write to the start of each chunk file.
    delete_on_error : `bool`, optional, default: `True`
        Delete all the chunk files if the chunking routine ends in an error.
    reverse : `bool`, optional, default: `False`
        Sort in reverse.
    chunk_prefix : `str`, optional, default: `NoneType`
        An optional prefix to add to the start of each chunk file name.
    chunk_suffix : `str`, optional, default: `.chunk.gz`
        An optional suffix to add to to the end of each chunk filename. Set to
        ``NoneType`` if you do not want a suffix.
    **write_kwargs
        Any keyword arguments that will be passed onto the ``write_method``.
    **csv_kwargs
        Any keyword arguments that will be passed onto csv.writer.

    Notes
    -----
    This will keep adding chunks to a chunk file if the sort order is
    maintained between the last element of the previous chunk and the first
    element of the current chunk. This can reduce the number of chunk files
    that are needed and is useful if you suspect that the input may be sorted
    already.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, *args, **kwargs):
        self.initialise(kwargs)
        super().__init__(*args, **kwargs)
