# -*- coding: utf-8 -*-
# Copyright (C) Duncan Macleod (2013)
#
# This file is part of GWpy.
#
# GWpy is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GWpy is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GWpy.  If not, see <http://www.gnu.org/licenses/>.

"""Input/Output utilities for LAL Cache files.
"""

from __future__ import division

import os.path
from math import ceil
from multiprocessing import (cpu_count, Process, Queue as ProcessQueue)
from six import string_types
import tempfile
import warnings
from gzip import GzipFile

from six import string_types

from numpy import recarray
from numpy.lib import recfunctions

try:
    from glue.lal import Cache
except ImportError:  # no lal
    HAS_CACHE = False
else:
    HAS_CACHE = True
    from lal.utils import CacheEntry

from astropy.table import (Table, vstack as vstack_tables)
from astropy.io.registry import _get_valid_format

from ..time import LIGOTimeGPS

__author__ = 'Duncan Macleod <duncan.macleod@ligo.org>'

# build list of file-like types
try:  # python2.x
    FILE_LIKE = [file, GzipFile]
except NameError:  # python3.x
    from io import IOBase
    FILE_LIKE = [IOBase, GzipFile]
try:  # protect against private member being removed
    FILE_LIKE.append(tempfile._TemporaryFileWrapper)
except AttributeError:
    pass
FILE_LIKE = tuple(FILE_LIKE)


def open_cache(lcf, coltype=LIGOTimeGPS):
    """Read a LAL-format cache file into memory as a `Cache`

    Parameters
    ----------
    lcf : `str`, `file`
        input file or file path to read

    coltype : `LIGOTimeGPS`, `int`, optional
        `type` for GPS times

    Returns
    -------
    cache : :class:`glue.lal.Cache`
        a cache object, representing each line in the file as a
        :class:`~lal.utils.CacheEntry`
    """
    from glue.lal import Cache
    if isinstance(lcf, FILE_LIKE):
        return Cache.fromfile(lcf, coltype=coltype)
    else:
        with open(lcf, 'r') as f:
            return Cache.fromfile(f, coltype=coltype)


def file_list(flist):
    """Parse a number of possible input types into a list of filepaths.

    Parameters
    ----------
    flist : `file-like` or `list-like` iterable
        the input data container, normally just a single file path, or a list
        of paths, but can generally be any of the following

        - `str` representing a single file path (or comma-separated collection)
        - open `file` or `~gzip.GzipFile` object
        - `~lal.utils.CacheEntry`
        - `~glue.lal.Cache` object or `str` with '.cache' or '.lcf extension
        - simple `list` or `tuple` of `str` paths

    Returns
    -------
    files : `list`
        `list` of `str` file paths

    Raises
    ------
    ValueError
        if the input `flist` cannot be interpreted as any of the above inputs
    """
    # detect open files
    if isinstance(flist, FILE_LIKE):
        flist = flist.name
    # then format list of file paths
    if HAS_CACHE and isinstance(flist, CacheEntry):
        return [flist.path]
    elif (isinstance(flist, string_types) and
          flist.endswith(('.cache', '.lcf'))):
        return open_cache(flist).pfnlist()
    elif isinstance(flist, string_types):
        return flist.split(',')
    elif HAS_CACHE and isinstance(flist, Cache):
        return flist.pfnlist()
    elif isinstance(flist, (list, tuple)):
        return flist
    raise ValueError("Could not parse input %r as one or more "
                     "file-like objects" % flist)


def file_segment(filename):
    """Return the data segment for a filename following T050017

    Parameters
    ---------
    filename : `str`, `~lal.utils.CacheEntry`
        the path name of a file

    Returns
    -------
    segment : `~gwpy.segments.Segment`
        the ``[start, stop)`` GPS segment covered by the given file

    Notes
    -----
    `LIGO-T050017 <https://dcc.ligo.org/LIGO-T050017/public>`_ declares
    a filenaming convention that includes documenting the GPS start integer
    and integer duration of a file, see that document for more details.
    """
    try:  # filename object provides its own segment information
        return filename.segment
    except AttributeError:  # otherwise parse from T050017 spec
        from ..segments import Segment
        base = os.path.basename(filename)
        _, _, s, e = base.split('-')
        s = int(s)
        e = int(e.split('.')[0])
        return Segment(s, s+e)


# ----------------------------------------------------------------------------
# generic multiprocessing from caches

def read_cache(cache, target, nproc, post, *args, **kwargs):
    """Read arbitrary data from a cache file

    Parameters
    ----------
    cache : :class:`glue.lal.Cache`, `str`
        cache of files files, or path to a LAL-format cache file
        on disk.
    target : `type`
        target class to read into.
    nproc : `int`
        number of individual processes to use.
    post : `function`
        function to post-process output object before returning.
        The output of this method will be returns, so in-place operations
        must return the object.
    *args
        other positional arguments to pass to the target.read()
        classmethod.
    **kwargs
        keyword arguments to pass to the target.read() classmethod.

    Returns
    -------
    data : target
        an instance of the target class, seeded with data read from
        the cache.

    Notes
    -----
    The returned object is constructed from the output of each
    sub-process via the '+=' in-place addition operator.

    If the input cache is indeed a :class:`~glue.lal.Cache` object,
    the sub-processes will be combined in time order, otherwise the ordering
    is given by the order of entries in the input cache (for example,
    if it is a simple `list` of files).

    .. warning::

       no protection is given against overloading the host, for example,
       no checks are done to ensure that ``nproc`` is less than the number
       of available cores.

       High values of ``nproc`` should be used at the users discretion,
       the GWpy team accepts to liability for loss as a result of abuse
       of this feature.
    """
    # read the cache
    if isinstance(cache, FILE_LIKE + string_types):
        cache = open_cache(cache)
    if HAS_CACHE and isinstance(cache, Cache):
        cache.sort(key=lambda ce: ce.segment[0])

    # force one file per process minimum
    nproc = min(nproc, len(cache))
    if nproc > cpu_count():
        warnings.warn("Using %d processes on a %d-core machine is "
                      "unrecommended...but not forbidden."
                      % (nproc, cpu_count()))

    # work out underlying data type
    try:
        kwargs.setdefault(
            'format', _get_valid_format('read', target, None,
                                        None, (cache[0],), {}))
    # if empty, put anything, since it doesn't matter
    except IndexError:
        kwargs.setdefault('format', 'ascii')
    except Exception:
        if 'format' not in kwargs:
            raise

    if nproc <= 1:
        return target.read(cache, *args, **kwargs)

    # define how to read each sub-cache
    def _read(q, sc, i):
        try:
            q.put((i, target.read(sc, *args, **kwargs)))
        except Exception as e:
            q.put(e)

    # separate cache into parts
    fperproc = int(ceil(len(cache) / nproc))
    subcaches = [cache.__class__(cache[i:i+fperproc]) for
                 i in range(0, len(cache), fperproc)]

    # start all processes
    queue = ProcessQueue(nproc)
    proclist = []
    for i, subcache in enumerate(subcaches):
        if len(subcache) == 0:
            continue
        process = Process(target=_read, args=(queue, subcache, i))
        process.daemon = True
        proclist.append(process)
        process.start()

    # get data and block
    pout = []
    for i in range(len(proclist)):
        result = queue.get()
        if isinstance(result, Exception):
            raise result
        pout.append(result)
    for process in proclist:
        process.join()

    # combine and return
    data = zip(*sorted(pout, key=lambda out: out[0]))[1]
    if issubclass(target, Table):  # astropy.table.Table
        out = vstack_tables(data, join_type='exact')
    elif issubclass(target, recarray):
        out = recfunctions.stack_arrays(data, asrecarray=True, usemask=False,
                                        autoconvert=True).view(target)
    else:
        try:
            if hasattr(target, 'tableName'):  # glue.ligolw.table.Table
                out = data[0]
            else:
                out = data[0].copy()
        except AttributeError:
            out = data[0]
        for datum in data[1:]:
            out += datum

    if post:
        return post(out)
    else:
        return out


def read_cache_factory(target):
    """Generate a read_cache method specific to a given target class.

    Parameters
    ----------
    target : `type`
        `class` object to bind to returned method

    Returns
    -------
    read_target_cache : `function`
        newly-minted method taking arbitrary arguments and keyword
        arguments, and handling type associations and multiprocessing.
    """
    def _read(f, *args, **kwargs):
        nproc = kwargs.pop('nproc', 1)
        post = kwargs.pop('post', None)
        return read_cache(f, target, nproc, post, *args, **kwargs)
    return _read


def cache_segments(*caches):
    """Build a `SegmentList` of data availability for these `Caches`

    Parameters
    ----------
    *cache : `~glue.lal.Cache`, `list`
        one of more `Cache` objects (or simple `list`)

    Returns
    -------
    segments : `~gwpy.segments.SegmentList`
        a list of segments for when data should be available
    """
    from ..segments import SegmentList
    out = SegmentList()
    for cache in caches:
        out.extend(file_segment(e) for e in cache)
    return out.coalesce()


def flatten(*caches):
    """Flatten a list of :class:`Caches <glue.lal.Cache>` into a single cache

    Parameters
    ----------
    *caches
        one or more :class:`~glue.lal.Cache` objects

    Returns
    -------
    flat : :class:`~glue.lal.Cache`
        a single cache containing the unique set of entries across
        each input
    """
    cache_type = type(caches[0])
    return cache_type([e for c in caches for e in c]).unique()


def find_contiguous(*caches):
    """Separate one or more caches into sets of contiguous caches

    Parameters
    ----------
    *caches
        one or more :class:`~glue.lal.Cache` objects

    Returns
    -------
    caches : `iter` of :class:`~glue.lal.Cache`
        an interable yielding each contiguous cache
    """
    from glue.lal import Cache
    try:
        flat = flatten(*caches)
    except IndexError:
        flat = Cache()
    for segment in cache_segments(flat):
        yield flat.sieve(segment=segment)
