# http://pyrocko.org - GPLv3
#
# The Pyrocko Developers, 21st Century
# ---|P------/S----------~Lg----------
from __future__ import division
import struct
import logging
import numpy as num

try:
    range = xrange
except NameError:
    pass

logger = logging.getLogger('pyrocko.spit')

or_ = num.logical_or
and_ = num.logical_and
not_ = num.logical_not
all_ = num.all
any_ = num.any


class OutOfBounds(Exception):
    pass


class Cell(object):
    def __init__(self, tree, index, f=None):
        self.tree = tree
        self.index = index
        self.depths = num.log2(index).astype(num.int)
        self.bad = False
        self.children = []
        n = 2**self.depths
        i = self.index - n
        delta = (self.tree.xbounds[:, 1] - self.tree.xbounds[:, 0])/n
        xmin = self.tree.xbounds[:, 0]
        self.xbounds = self.tree.xbounds.copy()
        self.xbounds[:, 0] = xmin + i * delta
        self.xbounds[:, 1] = xmin + (i+1) * delta
        self.a = self.xbounds[:, ::-1].copy()
        self.b = self.a.copy()
        self.b[:, 1] = self.xbounds[:, 1] - self.xbounds[:, 0]
        self.b[:, 0] = - self.b[:, 1]

        self.a[:, 0] += (self.b[:, 0] == 0.0)*0.5
        self.a[:, 1] -= (self.b[:, 1] == 0.0)*0.5
        self.b[:, 0] -= (self.b[:, 0] == 0.0)
        self.b[:, 1] += (self.b[:, 1] == 0.0)

        if f is None:
            it = nditer_outer(tuple(self.xbounds) + (None,))
            for vvv in it:
                vvv[-1][...] = self.tree._f_cached(vvv[:-1])

            self.f = it.operands[-1]
        else:
            self.f = f

    def interpolate(self, x):
        if self.children:
            for cell in self.children:
                if all_(and_(cell.xbounds[:, 0] <= x,
                        x <= cell.xbounds[:, 1])):
                    return cell.interpolate(x)

        else:
            if all_(num.isfinite(self.f)):
                ws = (x[:, num.newaxis] - self.a)/self.b
                wn = num.multiply.reduce(num.ix_(*ws))
                return num.sum(self.f * wn)
            else:
                return None

    def interpolate_many(self, x):
        if self.children:
            result = num.empty(x.shape[0], dtype=num.float)
            result[:] = None
            for cell in self.children:
                indices = num.where(
                    self.tree.ndim == num.sum(and_(
                        cell.xbounds[:, 0] <= x,
                        x <= cell.xbounds[:, 1]), axis=-1))[0]

                if indices.size != 0:
                    result[indices] = cell.interpolate_many(x[indices])

            return result

        else:
            if all_(num.isfinite(self.f)):
                ws = (x[..., num.newaxis] - self.a)/self.b
                npoints = ws.shape[0]
                ndim = self.tree.ndim
                ws_pimped = [ws[:, i, :] for i in range(ndim)]
                for i in range(ndim):
                    s = [npoints] + [1] * ndim
                    s[1+i] = 2
                    ws_pimped[i].shape = tuple(s)

                wn = ws_pimped[0]
                for idim in range(1, ndim):
                    wn = wn * ws_pimped[idim]

                result = wn * self.f
                for i in range(ndim):
                    result = num.sum(result, axis=-1)

                return result
            else:
                result = num.empty(x.shape[0], dtype=num.float)
                result[:] = None
                return result

    def slice(self, x):
        x = num.array(x, dtype=num.float)
        x_mask = not_(num.isfinite(x))
        x_ = x.copy()
        x_[x_mask] = 0.0
        return [
            cell for cell in self.children if all_(or_(
                x_mask,
                and_(
                    cell.xbounds[:, 0] <= x_,
                    x_ <= cell.xbounds[:, 1])))]

    def plot_rects(self, axes, x, dims):
        if self.children:
            for cell in self.slice(x):
                cell.plot_rects(axes, x, dims)

        else:
            points = []
            for iy, ix in ((0, 0), (0, 1), (1, 1), (1, 0), (0, 0)):
                points.append(
                    (self.xbounds[dims[0], iy], self.xbounds[dims[1], ix]))

            points = num.transpose(points)
            axes.plot(points[1], points[0], color=(0.1, 0.1, 0.0, 0.1))

    def check_holes(self):
        ''' Check if :py:class:`Cell` or its' children contain NaNs'''
        if self.children:
            return any([child.check_holes() for child in self.children])
        else:
            return num.any(num.isnan(self.f))

    def plot_2d(self, axes, x, dims):
        idims = num.array(dims)
        self.plot_rects(axes, x, dims)
        coords = [
            num.linspace(xb[0], xb[1], 1+int((xb[1]-xb[0])/d))
            for (xb, d) in zip(self.xbounds[idims, :], self.tree.xtols[idims])]

        npoints = coords[0].size * coords[1].size
        g = num.meshgrid(*coords[::-1])[::-1]
        points = num.empty((npoints, self.tree.ndim), dtype=num.float)
        for idim in range(self.tree.ndim):
            try:
                idimout = dims.index(idim)
                points[:, idim] = g[idimout].ravel()
            except ValueError:
                points[:, idim] = x[idim]

        fi = num.empty((coords[0].size, coords[1].size), dtype=num.float)
        fi_r = fi.ravel()
        fi_r[...] = self.interpolate_many(points)

        if num.any(num.isnan(fi)):
            logger.warn('')
        if any_(num.isfinite(fi)):
            fi = num.ma.masked_invalid(fi)
            axes.imshow(
                fi, origin='lower',
                extent=[coords[1].min(), coords[1].max(),
                        coords[0].min(), coords[0].max()],
                interpolation='nearest',
                aspect='auto',
                cmap='RdYlBu')

    def plot_1d(self, axes, x, dim):
        xb = self.xbounds[dim]
        d = self.tree.xtols[dim]
        coords = num.linspace(xb[0], xb[1], 1+int((xb[1]-xb[0])/d))

        npoints = coords.size
        points = num.empty((npoints, self.tree.ndim), dtype=num.float)
        for idim in range(self.tree.ndim):
            if idim == dim:
                points[:, idim] = coords
            else:
                points[:, idim] = x[idim]

        fi = self.interpolate_many(points)
        if any_(num.isfinite(fi)):
            fi = num.ma.masked_invalid(fi)
            axes.plot(coords, fi)

    def __iter__(self):
        yield self
        for c in self.children:
            for x in c:
                yield x

    def dump(self, file):
        self.index.astype('<i4').tofile(file)
        self.f.astype('<f8').tofile(file)
        for c in self.children:
            c.dump(file)


def bread(f, fmt):
    s = f.read(struct.calcsize(fmt))
    return struct.unpack(fmt, s)


class SPTree(object):

    def __init__(self, f=None, ftol=None, xbounds=None, xtols=None,
                 filename=None, addargs=()):

        '''Create n-dimensional space partitioning interpolator.

        :param f: callable function f(x) where x is a vector of size n
        :param ftol: target accuracy |f_interp(x) - f(x)| <= ftol
        :param xbounds: bounds of x, shape (n, 2)
        :param xtols: target coarsenesses in x, vector of size n
        :param addargs: additional arguments to pass to f
        '''

        if filename is None:
            assert all(v is not None for v in (f, ftol, xbounds, xtols))

            self.f = f
            self.ftol = float(ftol)
            self.f_values = {}
            self.ncells = 0
            self.addargs = addargs

            self.xbounds = num.asarray(xbounds, dtype=num.float)
            assert self.xbounds.ndim == 2
            assert self.xbounds.shape[1] == 2
            self.ndim = self.xbounds.shape[0]

            self.xtols = num.asarray(xtols, dtype=num.float)
            assert self.xtols.ndim == 1 and self.xtols.size == self.ndim

            self.maxdepths = num.ceil(num.log2(
                num.maximum(
                    1.0,
                    (self.xbounds[:, 1] - self.xbounds[:, 0]) / self.xtols)
                )).astype(num.int)

            self.root = None
            self.ones_int = num.ones(self.ndim, dtype=num.int)

            cc = num.ix_(*[num.arange(3)]*self.ndim)
            w = num.zeros([3]*self.ndim + [self.ndim, 2])
            for i, c in enumerate(cc):
                w[..., i, 0] = (2-c)*0.5
                w[..., i, 1] = c*0.5

            self.pointmaker = w
            self.pointmaker_mask = num.sum(w[..., 0] == 0.5, axis=-1) != 0
            self.pointmaker_masked = w[self.pointmaker_mask]

            self.nothing_found_yet = True

            self.root = Cell(self, self.ones_int)
            self.ncells += 1

            self.fraction_bad = 0.0
            self.nbad = 0
            self.cells_to_continue = []
            for clipdepth in range(0, num.max(self.maxdepths)+1):
                self.clipdepth = clipdepth
                self.tested = 0
                if self.clipdepth == 0:
                    self._fill(self.root)
                else:
                    self._continue_fill()

                self.status()

                if not self.cells_to_continue:
                    break

        else:
            self._load(filename)

    def status(self):
        perc = (1.0-self.fraction_bad)*100
        s = '%6.1f%%' % perc

        if self.fraction_bad != 0.0 and s == ' 100.0%':
            s = '~100.0%'

        logger.info('at level %2i: %s covered, %6i cell%s' % (
            self.clipdepth, s, self.ncells, ['s', ''][self.ncells == 1]))

    def __iter__(self):
        return iter(self.root)

    def __len__(self):
        return self.ncells

    def dump(self, filename):
        with open(filename, 'wb') as file:
            version = 1
            file.write(b'SPITREE ')
            file.write(struct.pack(
                '<QQQd', version, self.ndim, self.ncells, self.ftol))
            self.xbounds.astype('<f8').tofile(file)
            self.xtols.astype('<f8').tofile(file)
            self.root.dump(file)

    def _load(self, filename):
        with open(filename, 'rb') as file:
            marker, version, self.ndim, self.ncells, self.ftol = bread(
                file, '<8sQQQd')
            assert marker == b'SPITREE '
            assert version == 1
            self.xbounds = num.fromfile(
                file, dtype='<f8', count=self.ndim*2).reshape(self.ndim, 2)
            self.xtols = num.fromfile(
                file, dtype='<f8', count=self.ndim)

            path = []
            for icell in range(self.ncells):
                index = num.fromfile(
                    file, dtype='<i4', count=self.ndim)
                f = num.fromfile(
                    file, dtype='<f8', count=2**self.ndim).reshape(
                        [2]*self.ndim)

                cell = Cell(self, index, f)
                if not path:
                    self.root = cell
                    path.append(cell)

                else:
                    while not any_(path[-1].index == (cell.index >> 1)):
                        path.pop()

                    path[-1].children.append(cell)
                    path.append(cell)

    def _f_cached(self, x):
        return getset(
            self.f_values, tuple(float(xx) for xx in x), self.f, self.addargs)

    def interpolate(self, x):
        x = num.asarray(x, dtype=num.float)
        assert x.ndim == 1 and x.size == self.ndim
        if not all_(and_(self.xbounds[:, 0] <= x, x <= self.xbounds[:, 1])):
            raise OutOfBounds()

        return self.root.interpolate(x)

    def __call__(self, x):
        return self.interpolate(x)

    def interpolate_many(self, x):
        return self.root.interpolate_many(x)

    def _continue_fill(self):
        cells_to_continue, self.cells_to_continue = self.cells_to_continue, []
        for cell in cells_to_continue:
            self._deepen_cell(cell)

    def _fill(self, cell):

        self.tested += 1
        xtestpoints = num.sum(cell.xbounds * self.pointmaker_masked, axis=-1)

        fis = cell.interpolate_many(xtestpoints)
        fes = num.array(
            [self._f_cached(x) for x in xtestpoints], dtype=num.float)

        iffes = num.isfinite(fes)
        iffis = num.isfinite(fis)
        works = iffes == iffis
        iif = num.logical_and(iffes, iffis)

        works[iif] *= num.abs(fes[iif] - fis[iif]) < self.ftol

        nundef = num.sum(not_(num.isfinite(fes))) + \
            num.sum(not_(num.isfinite(cell.f)))

        some_undef = 0 < nundef < (xtestpoints.shape[0] + cell.f.size)

        if any_(works):
            self.nothing_found_yet = False

        if not all_(works) or some_undef or self.nothing_found_yet:
            deepen = self.ones_int.copy()
            if not some_undef:
                works_full = num.ones([3]*self.ndim, dtype=num.bool)
                works_full[self.pointmaker_mask] = works
                for idim in range(self.ndim):
                    dimcorners = [slice(None, None, 2)] * self.ndim
                    dimcorners[idim] = 1
                    if all_(works_full[tuple(dimcorners)]):
                        deepen[idim] = 0

            if not any_(deepen):
                deepen = self.ones_int

            deepen = num.where(
                cell.depths + deepen > self.maxdepths, 0, deepen)

            cell.deepen = deepen

            if any_(deepen) and all_(cell.depths + deepen <= self.clipdepth):
                self._deepen_cell(cell)
            else:
                if any_(deepen):
                    self.cells_to_continue.append(cell)

                cell.bad = True
                self.fraction_bad += num.product(1.0/2**cell.depths)
                self.nbad += 1

    def _deepen_cell(self, cell):
        if cell.bad:
            self.fraction_bad -= num.product(1.0/2**cell.depths)
            self.nbad -= 1
            cell.bad = False

        for iadd in num.ndindex(*(cell.deepen+1)):
            index_child = (cell.index << cell.deepen) + iadd
            child = Cell(self, index_child)
            self.ncells += 1
            cell.children.append(child)
            self._fill(child)

    def check_holes(self):
        '''Check for NaNs in :py:class:`SPTree`'''
        return self.root.check_holes()

    def plot_2d(self, axes=None, x=None, dims=None):
        assert self.ndim >= 2

        if x is None:
            x = num.zeros(self.ndim)
            x[-2:] = None

        x = num.asarray(x, dtype=num.float)
        if dims is None:
            dims = [i for (i, v) in enumerate(x) if not num.isfinite(v)]

        assert len(dims) == 2

        plt = None
        if axes is None:
            from matplotlib import pyplot as plt
            axes = plt.gca()

        self.root.plot_2d(axes, x, dims)

        axes.set_xlabel('Dim %i' % dims[1])
        axes.set_ylabel('Dim %i' % dims[0])

        if plt:
            plt.show()

    def plot_1d(self, axes=None, x=None, dims=None):

        if x is None:
            x = num.zeros(self.ndim)
            x[-1:] = None

        x = num.asarray(x, dtype=num.float)
        if dims is None:
            dims = [i for (i, v) in enumerate(x) if not num.isfinite(v)]

        assert len(dims) == 1

        plt = None
        if axes is None:
            from matplotlib import pyplot as plt
            axes = plt.gca()

        self.root.plot_1d(axes, x, dims[0])

        axes.set_xlabel('Dim %i' % dims[0])

        if plt:
            plt.show()


def getset(d, k, f, addargs):
    try:
        return d[k]
    except KeyError:
        v = d[k] = f(k, *addargs)
        return v


def nditer_outer(x):
    add = []
    if x[-1] is None:
        x_ = x[:-1]
        add = [None]
    else:
        x_ = x

    return num.nditer(
        x,
        op_axes=(num.identity(len(x_), dtype=num.int)-1).tolist() + add)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)

    def f(x):
        x0 = num.array([0.5, 0.5, 0.5])
        r = 0.5
        if num.sqrt(num.sum((x-x0)**2)) < r:

            return x[2]**4 + x[1]

        return None

    tree = SPTree(f, 0.01, [[0., 1.], [0., 1.], [0., 1.]], [0.025, 0.05, 0.1])

    import tempfile
    import os
    fid, fn = tempfile.mkstemp()
    tree.dump(fn)
    tree = SPTree(filename=fn)
    os.unlink(fn)

    from matplotlib import pyplot as plt

    v = 0.5
    axes = plt.subplot(2, 2, 1)
    tree.plot_2d(axes, x=(v, None, None))
    axes = plt.subplot(2, 2, 2)
    tree.plot_2d(axes, x=(None, v, None))
    axes = plt.subplot(2, 2, 3)
    tree.plot_2d(axes, x=(None, None, v))

    axes = plt.subplot(2, 2, 4)
    tree.plot_1d(axes, x=(v, v, None))

    plt.show()
