# cython: c_string_type=unicode

"""Readers and writers

The basic idea here is to provide Python IO semantics by using Python to do
the IO. Otherwise we get C++ IO semantics in Python. It also means we can use
dummy files, e.g. anything with read/write attributes. Generally a much better
idea than just 'give this a filename', and well worth the inefficiencies and
potential memory limits.
"""

import sys
from libcpp.unordered_map cimport unordered_map
from cython.operator cimport preincrement as inc

## Make a string that is a comma-separated list of patterns
def _prep_patterns(patterns):
    if not patterns:
        patterns = [ ]
    elif not isinstance(patterns, (list,tuple)):
        patterns = [patterns]
    return ','.join(patterns)

def _mapTypes(cpptype):
    if 'Binned' not in cpptype:
        return cpptype
    pos = cpptype.find('<')
    base = cpptype[6:pos] # remove leading 'Binned' and axes tuple
    N = str(len(cpptype[pos+1:-1].split(','))) # count axes
    return base+N+'D'

## Make a Python list of analysis objects from a C++ vector of them
cdef list _aobjects_to_list(vector[c.AnalysisObject*]* aobjects):
    cdef list out = []
    cdef c.AnalysisObject* ao
    cdef size_t i
    for i in range(aobjects.size()):
        ao = deref(aobjects)[i]
        ## NOTE: automatic type conversion by passing the type() as a key to globals()
        pytype = _mapTypes(ao.type())
        newao = cutil.new_owned_cls(globals()[pytype], ao)
        ## try setting axis configuration for binned objects
        try:
            newao._set_config()
        except:
            pass
        out.append(newao)
    return out

## Make a Python dict of analysis objects from a C++ vector of them
cdef dict _aobjects_to_dict(vector[c.AnalysisObject*]* aobjects):
    cdef dict out = {}
    # from collections import OrderedDict
    # OrderedDict out = OrderedDict()
    cdef c.AnalysisObject* ao
    cdef size_t i
    for i in range(aobjects.size()):
        ao = deref(aobjects)[i]
        ## NOTE: automatic type conversion by passing the type() as a key to globals()
        pytype = _mapTypes(ao.type())
        newao = cutil.new_owned_cls( globals()[pytype], ao)
        ## try setting axis configuration for binned objects
        try:
            newao._set_config()
        except:
            pass
        out[newao.path()] = newao
    return out

## Make a Python dict of dicts for file index from a C++ unordered_map of unordered_maps
cdef dict _idxMap_to_dict(unordered_map[string, unordered_map[string, size_t]]* idxMap):
    cdef dict out = {}
    cdef unordered_map[string, unordered_map[string, size_t]].iterator it = idxMap.begin()
    cdef dict innerDict = {}
    cdef unordered_map[string, size_t].iterator it_inner

    while(it != idxMap.end()):
        objType = deref(it).first
        nestedMap = deref(it).second
        innerDict = {}
        it_inner = nestedMap.begin()

        while(it_inner != nestedMap.end()):
            path = deref(it_inner).first
            binNum = deref(it_inner).second
            innerDict[path] = binNum
            inc(it_inner)

        out[objType] = innerDict
        inc(it)

    return out


def _bytestr_from_file(file_or_filename):
    """\
    Read a file's contents as a returned bytestring.

    The file argument can either be a file object, filename, or
    special "-" reference to stdin.
    """
    if hasattr(file_or_filename, "read"):
        s = file_or_filename.read()
    elif file_or_filename == "-":
        s = sys.stdin.buffer.read()
    else:
        with open(file_or_filename, "rb") as f:
            s = f.read()
    if type(s) is str:
        s = s.encode("utf-8")
    return s

def _str_to_file(s, file_or_filename):
    """\
    Write a string to a file.

    The file argument can either be a file object, filename, or
    special "-" reference to stdout.
    """
    if hasattr(file_or_filename, "write"):
        file_or_filename.write(s)
    elif file_or_filename == "-":
        sys.stdout.write(s)
    else:
        with open(file_or_filename, "w") as f:
            f.write(s)

def _istxt(x):
    if sys.version_info < (3, 0):
        return type(x) in [unicode, str]
    else:
        return type(x) is str

def _mktxtifstr(x):
    if sys.version_info < (3, 0):
        return unicode(x) if type(x) is str else x
    return x



##
## Readers
##

def read(filename, asdict=True, patterns='', unpatterns=''):
    """
    Read data objects from the provided filename, auto-determining the format
    from the file extension.

    The loaded data objects can be filtered on their path strings, using the
    optional patterns and unpatterns arguments. These can be strings, compiled
    regex objects with a 'match' method, or any iterable of those types. If
    given, only analyses with paths which match at least one pattern, and do not
    match any unpatterns, will be returned.

    Returns a dict or list of analysis objects depending on the asdict argument.
    """
    cdef vector[c.AnalysisObject*] aobjects
    cdef c.istringstream iss
    filename = _mktxtifstr(filename)
    patterns = _prep_patterns(patterns)
    unpatterns = _prep_patterns(unpatterns)
    if _istxt(filename):
        c.IO_read_from_file(filename, aobjects, patterns, unpatterns)
    else:
        s = _bytestr_from_file(filename)
        iss.str(s)
        c.IO_read_from_stringstream(iss, aobjects, 'yoda', patterns, unpatterns)

    if asdict:
        d = _aobjects_to_dict(&aobjects)
        try:
            import collections
            return collections.OrderedDict(sorted(d.items()))
        except:
            return d
    else:
        return _aobjects_to_list(&aobjects)


def readYODA(file_or_filename, asdict=True, patterns='', unpatterns=''):
    """
    Read data objects from the provided YODA-format file.

    The loaded data objects can be filtered on their path strings, using the
    optional patterns and unpatterns arguments. These can be strings, compiled
    regex objects with a 'match' method, or any iterable of those types. If
    given, only analyses with paths which match at least one pattern, and do not
    match any unpatterns, will be returned.

    Returns a dict or list of analysis objects depending on the asdict argument.
    """
    cdef c.istringstream iss
    cdef vector[c.AnalysisObject*] aobjects
    s = _bytestr_from_file(file_or_filename)
    iss.str(s)
    patterns = _prep_patterns(patterns)
    unpatterns = _prep_patterns(unpatterns)
    c.ReaderYODA_create().read(iss, aobjects, patterns, unpatterns)
    return _aobjects_to_dict(&aobjects) if asdict \
        else _aobjects_to_list(&aobjects)


def readFLAT(file_or_filename, asdict=True, patterns='', unpatterns=''):
    """
    Read data objects from the provided FLAT-format file.

    The loaded data objects can be filtered on their path strings, using the
    optional patterns and unpatterns arguments. These can be strings, compiled
    regex objects with a 'match' method, or any iterable of those types. If
    given, only analyses with paths which match at least one pattern, and do not
    match any unpatterns, will be returned.

    Returns a dict or list of analysis objects depending on the asdict argument.
    """
    cdef c.istringstream iss
    cdef vector[c.AnalysisObject*] aobjects
    s = _bytestr_from_file(file_or_filename)
    iss.str(s)
    patterns = _prep_patterns(patterns)
    unpatterns = _prep_patterns(unpatterns)
    c.ReaderFLAT_create().read(iss, aobjects, patterns, unpatterns)
    return _aobjects_to_dict(&aobjects) if asdict \
        else _aobjects_to_list(&aobjects)


##
## Writers
##

def write(ana_objs, filename, precision=-1):
    """
    Write data objects to the provided filename,
    auto-determining the format from the file extension.
    """
    # cdef c.ostringstream oss
    cdef vector[c.AnalysisObject*] vec
    cdef AnalysisObject a
    aolist = [ao for key,ao in sorted(ana_objs.items())] if hasattr(ana_objs, "items") \
              else ana_objs if hasattr(ana_objs, "__iter__") else [ana_objs]
    for a in aolist:
        vec.push_back(a.aoptr())
    c.IO_write_to_file(filename, vec, precision)
    #_str_to_file(oss.str(), filename)


def writeYODA(ana_objs, file_or_filename, precision=-1):
    """
    Write data objects to the provided file in YODA format.
    """
    cdef c.ostringstream oss
    cdef c.Writer* w
    cdef vector[c.AnalysisObject*] vec
    cdef AnalysisObject a
    aolist = list(ana_objs.values()) if hasattr(ana_objs, "values") else [ana_objs] \
             if issubclass(type(ana_objs), AnalysisObject) else ana_objs
    for a in aolist:
        vec.push_back(a.aoptr())
    w = & c.WriterYODA_create()
    if precision >= 0:
        w.setPrecision(precision)
    if _istxt(file_or_filename):
        w.write_to_file(file_or_filename, vec)
    else:
        w.write(oss, vec)
        _str_to_file(oss.str(), file_or_filename)


def writeYODA1(ana_objs, file_or_filename, precision=-1):
    """
    Write data objects to the provided file in YODA1 format.
    """
    cdef c.ostringstream oss
    cdef c.Writer* w
    cdef vector[c.AnalysisObject*] vec
    cdef AnalysisObject a
    aolist = list(ana_objs.values()) if hasattr(ana_objs, "values") else ana_objs \
             if hasattr(ana_objs, "__iter__") else [ana_objs]
    for a in aolist:
        vec.push_back(a.aoptr())
    w = & c.WriterYODA1_create()
    if precision >= 0:
        w.setPrecision(precision)
    if _istxt(file_or_filename):
        w.write_to_file(file_or_filename, vec)
    else:
        w.write(oss, vec)
        _str_to_file(oss.str(), file_or_filename)


def writeFLAT(ana_objs, file_or_filename, precision=-1):
    """
    Write data objects to the provided file in FLAT format.
    """
    cdef c.ostringstream oss
    cdef c.Writer* w
    cdef vector[c.AnalysisObject*] vec
    cdef AnalysisObject a
    aolist = list(ana_objs.values()) if hasattr(ana_objs, "values") else ana_objs \
             if hasattr(ana_objs, "__iter__") else [ana_objs]
    for a in aolist:
        vec.push_back(a.aoptr())
    w = & c.WriterFLAT_create()
    if precision >= 0:
        w.setPrecision(precision)
    if _istxt(file_or_filename):
        w.write_to_file(file_or_filename, vec)
    else:
        w.write(oss, vec)
        _str_to_file(oss.str(), file_or_filename)

