"""
Fast arithmetic with simplicial univariate rational functions.
"""

from sage.all import *

from collections import namedtuple

from array import array

import ctypes, gzip, struct

import common

from util import create_logger, TemporaryDirectory, readable_filesize
logger = create_logger(__name__)

#
# Try to load the C cruncher. If that fails, we'll use a Python implementation.
#
from pkg_resources import resource_exists, resource_filename
try:
    if resource_exists('zetalib', 'crunch.so'):
        common.libcrunch = ctypes.CDLL(resource_filename('zetalib', 'crunch.so'))
except Exception as e:
    logger.warn('The C extension crunch.so could not be loaded. '
                'Computations of topological zeta functions will use '
                'the slower Python implementation.')
    # To debug uncomment:
    #logger.warn(e)

# [S]implicial [U]nivariate [R]ational [F]unctions are expressions of the
# form scalar / prod(rays[i][0]*s - rays[i][1], i=1,...,e).

SURF = namedtuple('SURF', ['scalar', 'rays'])
class SURFError(Exception):
    pass

_BUFSIZE = 32*1024*1024
INTSIZE = 4

if array('i').itemsize != INTSIZE:
    raise SURFError('Need sizeof(int) == 4')

class SURFSum(object):
    def __init__(self, filename, compresslevel=9):
        self._filename = filename
        self._meta_filename = filename + '.meta'
        self._compresslevel = compresslevel

        try:
            f = open(self._meta_filename, 'rb')
        except IOError:
            # No metadata -> start from scratch.
            self._critical = set()
            self._cand = dict()
            self._count = 0
            self._file = gzip.open(filename, 'wb', compresslevel)
        else:
            # Recover metadata and open file for appending.
            with f:
                D = loads(f.read())
            self._cand, self._critical, self._count = D['cand'], D['critical'], D['count']
            self._file = gzip.open(filename, 'ab', compresslevel)

    def __str__(self):
        return 'Sum of %d SURFs [%d critical points]' % (self._count, len(self._critical))

    def close(self):
        # Flush data file and write metadata.
        if not self._file.closed:
            self._file.close()
        with open(self._meta_filename, 'wb') as f:
            f.write(dumps(
                    { 'cand':     self._cand,
                      'critical': self._critical,
                      'count':    self._count }
                    ))

    def add(self, S):
        if self._file.closed:
            raise SURFError('Lost access to the SURFSum file')

        self._count += 1

        # Binary format:
        # k scalar a[0] b[0] a[1] b[1] ... a[k-1] b[k-1]

        # NOTE:
        # In rare cases, we can run into numbers that don't fit into 32-bit
        # integers. We then use an ad hoc hack to reduce to smaller numbers.
        # To trigger this, count ideals in L(6,11).

        raw = map(int, [len(S.rays), S.scalar])
        for (a,b) in S.rays:
            try:
                _ = array('i', [int(a), int(b)])
                raw.extend([int(a), int(b)])
            except OverflowError:
                if a:
                    raise SURFError('Ray does not fit into a pair of 32-bit integers')
            else:
                continue

            fac = factor(b)
            li = flatten([[p]*e for (p,e) in fac])
            li[0] *= fac.unit()
            # assert b == prod(li)
            if len(li) % 2 == 0:
                li[0] = -li[0]
            # assert -b == prod(-c for c in li)
            for c in li:
                raw.extend([int(0), int(c)])
            raw[0] += len(li)-1

        try:
            self._file.write(array('i', raw).tostring())
        except OverflowError:
            raise SURFError('Number too large to fit into a 32-bit integer')

        # Update the candidate denominator.
        E = {}
        for (a,b) in S.rays:
            if a == 0:
                continue

            self._critical.add(QQ(b)/QQ(a))

            # Only consider a*s-b with gcd(a,b) = 1 and a > 0.
            g = int(gcd((a, b)))
            a,b = a/g, b/g

            if a < 0:
                a,b = -a, -b

            # (possible, depending on the counting problem) TODO:
            # Get rid of things like s+1 (for subobjects) that cannot show up
            # in the final result.

            if E.has_key((a,b)):
                E[(a,b)] += 1
            else:
                E[(a,b)] = 1

        # Now 'E' gives the multiplicities of candidate terms for S.
        # We take the largest multiplicities over all 'S'.
        for r in E.keys():
            if (not self._cand.has_key(r)) or self._cand[r] < E[r]:
                self._cand[r] = E[r]

def _crunch_c(args):
    argc = int(len(args))
    argv = (ctypes.c_char_p*int(argc))()
    for i in xrange(argc):
        argv[i] = args[i]
    return common.libcrunch.crunch(argc,argv)

# A python implementation of `crunch.c'.
def _crunch_py(argv):

    if len(argv) < 4: return 1

    with open(argv[1], 'r') as file:
        nvalues = int(file.readline())
        values = [QQ(v) for v in file]

    if len(values) != nvalues: return 1

    results = nvalues * [QQ(0)]

    for dfile in argv[3:]:
        with gzip.open(dfile, 'rb') as gzfile:
            gzfile._read_eof = lambda *args: None # ignore gzip CRC errors in streams

            while True:
                raw = gzfile.read(INTSIZE)
                if not raw:
                    break # EOF

                if len(raw) != INTSIZE: return 1

                nrays, = struct.unpack('i', raw)
                chunk_size = 2*nrays + 1

                raw = gzfile.read(chunk_size * INTSIZE)
                if len(raw) != chunk_size * INTSIZE: return 1
                chunk = [QQ(v) for v in struct.unpack('i' * chunk_size, raw)]
                
                for i in xrange(nvalues):
                    results[i] += chunk[0] / prod(chunk[j]*values[i] - chunk[j+1] for j in xrange(1,chunk_size,2))

    with open(argv[2], 'w') as file:
        for r in results:
            file.write(str(r) + '\n')
    return 0

def crunch(args):
    return _crunch_c(args) if common.libcrunch else _crunch_py(args)

def multicrunch(surfsums, varname=None):
    """
    Given an iterable consisting of SURFSums, compute the rational function
    given by their combined sum.
    Note that this rational function necessarily has degree <= 0.
    """
    
    surfsums = list(surfsums)

    #
    # Combine the various critical sets and construct a candidate denominator.
    #

    critical = set().union(*(Q._critical for Q in surfsums))
    cand = dict()
    for Q in surfsums:
        E = Q._cand
        for r in E.keys():
            if not cand.has_key(r) or cand[r] < E[r]:
                cand[r] = E[r]

    if varname is None:
        varname = 's'

    R = QQ[varname]
    s = R.gen(0)
    g = R(prod((a*s-b)**e for ((a,b),e) in cand.iteritems()))
    m = g.degree()

    logger.info('Total number of SURFs: %d' % sum(Q._count for Q in surfsums))

    for Q in surfsums:
        Q._file.flush()

    logger.info('Combined size of data files: %s' %
         readable_filesize(sum(os.path.getsize(Q._filename) for Q in surfsums))
         )
    logger.info('Number of critical points: %d' % len(critical))
    logger.info('Degree of candidate denominator: %d' % m)

    #
    # Construct m + 1 non-critical points for evaluation.
    #

    values = set()
    while len(values) < m + 1:
        x = QQ.random_element()
        if x in critical:
            continue
        values.add(x)
    values = list(values)

    #
    # Set up parallel computations.
    #

    bucket_size = ceil(float(len(values))/common.ncpus)

    dat_filenames = [Q._filename for Q in surfsums]

    res_names = []
    val_names = []

    value_batches = [values[j::common.ncpus] for j in xrange(common.ncpus)]

    with TemporaryDirectory() as tmpdir:
        for j,v in enumerate(value_batches):
            if not v:
                break

            val_filename = os.path.join(tmpdir, 'values%d' % j)
            val_names.append(val_filename)
            res_names.append(os.path.join(tmpdir, 'results%d' % j))
            with open(val_filename, 'w') as val_file:
                val_file.write(str(len(v)) + '\n')
                for x in v:
                    val_file.write(str(x) + '\n')

        def fun(k):
            ret = crunch(['crunch', val_names[k], res_names[k]]+dat_filenames)
            if ret == 0:
                logger.info('Cruncher #%d finished.' % k)
            return ret

        fun = sage.parallel.decorate.parallel(ncpus=len(res_names))(fun)
        logger.info('Launching %d crunchers.' % len(res_names))

        for (arg, ret) in fun(range(len(res_names))):
            if ret == 'NO DATA':
                raise RuntimeError('A parallel process died')
            if ret != 0:
                raise RuntimeError('crunch failed')

        #
        # Collect results
        #
        pairs = []

        for j,rn in enumerate(res_names):
            it_batch = iter(value_batches[j])
            with open(rn, 'r') as res_file:
                for line in res_file:
                    # We also need to evaluate the candidate denominator 'g'
                    # from above at the given random points.
                    x = QQ(next(it_batch))
                    pairs.append((x, g(x)*QQ(line)))
            
    if len(values) != len(pairs):
        raise RuntimeError('Length of results is off') 

    f = R.lagrange_polynomial(list(pairs))
    res = SR(f/g)
    return res.factor() if res else res
