"""
A class hierarchy relating to fields of all kinds.
"""
from __future__ import print_function, division
import numpy as np
from ciabatta.meta import make_repr_str
from fealty import lattice, field_numerics, walled_field_numerics


class Space(object):

    def __init__(self, L, dim):
        self.L = L
        self.dim = dim

    @property
    def L_half(self):
        return self.L / 2.0

    @property
    def A(self):
        return self.L ** self.dim

    def iterate(self, *args, **kwargs):
        pass

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim)]
        return make_repr_str(self, fs)


class Field(Space):

    def __init__(self, L, dim, dx):
        Space.__init__(self, L, dim)
        self.M = int(round(self.L / dx))

    @property
    def dx(self):
        return self.L / self.M

    @property
    def A_i(self):
        return self.M ** self.dim

    @property
    def dA(self):
        return self.dx ** self.dim

    def density_field(self, r):
        return density(r, self.L, self.dx)

    def r_to_i(self, r):
        return lattice.r_to_i(r, self.L, self.dx)

    def i_to_r(self, i):
        return lattice.i_to_r(i, self.L, self.dx)

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim), ('dx', self.dx)]
        return make_repr_str(self, fs)


class Scalar(Field):

    def __init__(self, L, dim, dx, a_0=0.0):
        Field.__init__(self, L, dim, dx)
        self.a = np.ones(self.dim * (self.M,), dtype=np.float) * a_0

    def grad(self):
        return _grad(self.a, self.dx)

    def grad_i(self, r):
        return _grad_i(self.a, self.r_to_i(r), self.dx)

    def laplacian(self):
        return _laplace(self.a, self.dx)

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim), ('dx', self.dx),
              ('a_0', self.a_0)]
        return make_repr_str(self, fs)


class Diffusing(Scalar):

    def __init__(self, L, dim, dx, D, dt, a_0=0.0):
        Scalar.__init__(self, L, dim, dx, a_0=a_0)
        self.D = D
        self.dt = dt

        if self.D > self.dx ** 2 / (2.0 * self.dim * self.dt):
            raise Exception('Unstable diffusion constant')

    def iterate(self):
        self.a += self.D * self.laplacian() * self.dt

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim), ('dx', self.dx),
              ('D', self.D), ('dt', self.dt), ('a_0', self.a_0)]
        return make_repr_str(self, fs)


class WalledScalar(Scalar):

    def __init__(self, L, dim, dx, walls, a_0=0.0):
        Scalar.__init__(self, L, dim, dx, a_0=a_0)
        self.walls = walls
        # Make field zero-valued where obstructed
        self.a *= np.logical_not(self.walls)

    def grad(self):
        return _walled_grad(self.a, self.dx, self.walls)

    def grad_i(self, r):
        return _walled_grad_i(self.a, self.r_to_i(r), self.dx,
                              self.walls)

    def laplacian(self):
        return _walled_laplace(self.a, self.dx, self.walls)

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim), ('dx', self.dx),
              ('walls', self.walls), ('a_0', self.a_0)]
        return make_repr_str(self, fs)


# Note, inheritance order matters to get walled grad & laplacian call
# (see diamond problem on wikipedia and how python handles it)
class WalledDiffusing(WalledScalar, Diffusing):

    def __init__(self, L, dim, dx, walls, D, dt, a_0=0.0):
        Diffusing.__init__(self, L, dim, dx, D, dt, a_0=a_0)
        WalledScalar.__init__(self, L, dim, dx, walls, a_0=a_0)

    def __repr__(self):
        fs = [('L', self.L), ('dim', self.dim), ('dx', self.dx),
              ('walls', self.walls), ('D', self.D), ('dt', self.dt),
              ('a_0', self.a_0)]
        return make_repr_str(self, fs)


def density(r, L, dx):
    assert r.ndim == 2
    M = int(round(L / dx))
    dx = L / M
    inds = lattice.r_to_i(r, L, dx)
    f = np.zeros(r.shape[1] * (M,), dtype=np.int)
    if f.ndim == 1:
        field_numerics.density_1d(inds, f)
    elif f.ndim == 2:
        field_numerics.density_2d(inds, f)
    elif f.ndim == 3:
        field_numerics.density_3d(inds, f)
    else:
        raise Exception('Density calc not implemented in this dimension')
    return f / dx ** r.shape[1]


def _laplace(field, dx):
    assert dx > 0.0
    laplace = np.empty_like(field)
    if field.ndim == 1:
        field_numerics.laplace_1d(field, laplace, dx)
    elif field.ndim == 2:
        field_numerics.laplace_2d(field, laplace, dx)
    elif field.ndim == 3:
        field_numerics.laplace_3d(field, laplace, dx)
    else:
        raise Exception('Laplacian not implemented in this dimension')
    return laplace


def _grad_i(field, inds, dx):
    assert dx > 0.0
    assert inds.ndim == 2
    assert field.ndim == inds.shape[1]
    grad_i = np.empty(inds.shape, dtype=field.dtype)
    if field.ndim == 1:
        field_numerics.grad_i_1d(field, inds, grad_i, dx)
    elif field.ndim == 2:
        field_numerics.grad_i_2d(field, inds, grad_i, dx)
    elif field.ndim == 3:
        field_numerics.grad_i_3d(field, grad_i, dx)
    else:
        raise Exception("Grad_i not implemented in this dimension")
    return grad_i


def _grad(field, dx):
    assert dx > 0.0
    grad = np.empty(field.shape + (field.ndim,), dtype=field.dtype)
    if field.ndim == 1:
        field_numerics.grad_1d(field, grad, dx)
    elif field.ndim == 2:
        field_numerics.grad_2d(field, grad, dx)
    elif field.ndim == 3:
        field_numerics.grad_3d(field, grad, dx)
    else:
        raise Exception('Grad not implemented in this dimension')
    return grad


def _div(field, dx):
    assert dx > 0.0
    div = np.empty(field.shape[:-1], dtype=field.dtype)
    if field.ndim == 2:
        field_numerics.div_1d(field, div, dx)
    elif field.ndim == 3:
        field_numerics.div_2d(field, div, dx)
    elif field.ndim == 4:
        field_numerics.div_3d(field, div, dx)
    else:
        raise Exception('Divergence not implemented in this dimension')
    return div


def _walled_grad(field, dx, walls):
    assert field.shape == walls.shape
    assert dx > 0.0
    grad = np.empty(field.shape + (field.ndim,), dtype=field.dtype)
    if field.ndim == 1:
        walled_field_numerics.grad_1d(field, grad, dx, walls)
    elif field.ndim == 2:
        walled_field_numerics.grad_2d(field, grad, dx, walls)
    elif field.ndim == 3:
        walled_field_numerics.grad_3d(field, grad, dx, walls)
    else:
        raise Exception("Walled grad not implemented in this dimension")
    return grad


def _walled_grad_i(field, inds, dx, walls):
    assert field.shape == walls.shape
    assert dx > 0.0
    assert inds.ndim == 2
    assert field.ndim == inds.shape[1]
    grad_i = np.empty(inds.shape, dtype=field.dtype)
    if field.ndim == 1:
        walled_field_numerics.grad_i_1d(field, inds, grad_i, dx, walls)
    elif field.ndim == 2:
        walled_field_numerics.grad_i_2d(field, inds, grad_i, dx, walls)
    elif field.ndim == 3:
        walled_field_numerics.grad_i_3d(field, inds, grad_i, dx, walls)
    else:
        raise Exception("Walled Grad_i not implemented in this dimension")
    return grad_i


def _walled_laplace(field, dx, walls):
    assert field.shape == walls.shape
    assert dx > 0.0
    laplace = np.empty_like(field)
    if field.ndim == 1:
        walled_field_numerics.laplace_1d(field, laplace, dx, walls)
    elif field.ndim == 2:
        walled_field_numerics.laplace_2d(field, laplace, dx, walls)
    elif field.ndim == 3:
        walled_field_numerics.laplace_3d(field, laplace, dx, walls)
    else:
        raise Exception('Laplacian not implemented in this dimension')
    return laplace

[((147, 11, 147, 35), 'fealty.lattice.r_to_i', 'lattice.r_to_i', ({(147, 26, 147, 27): 'r', (147, 29, 147, 30): 'L', (147, 32, 147, 34): 'dx'}, {}), '(r, L, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((148, 8, 148, 49), 'numpy.zeros', 'np.zeros', (), '', True, 'import numpy as np\n'), ((162, 14, 162, 34), 'numpy.empty_like', 'np.empty_like', ({(162, 28, 162, 33): 'field'}, {}), '(field)', True, 'import numpy as np\n'), ((178, 13, 178, 52), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((192, 11, 192, 67), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((206, 10, 206, 55), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((221, 11, 221, 67), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((238, 13, 238, 52), 'numpy.empty', 'np.empty', (), '', True, 'import numpy as np\n'), ((253, 14, 253, 34), 'numpy.empty_like', 'np.empty_like', ({(253, 28, 253, 33): 'field'}, {}), '(field)', True, 'import numpy as np\n'), ((29, 15, 29, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(29, 29, 29, 33): 'self', (29, 35, 29, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((54, 15, 54, 49), 'fealty.lattice.r_to_i', 'lattice.r_to_i', ({(54, 30, 54, 31): 'r', (54, 33, 54, 39): 'self.L', (54, 41, 54, 48): 'self.dx'}, {}), '(r, self.L, self.dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((57, 15, 57, 49), 'fealty.lattice.i_to_r', 'lattice.i_to_r', ({(57, 30, 57, 31): 'i', (57, 33, 57, 39): 'self.L', (57, 41, 57, 48): 'self.dx'}, {}), '(i, self.L, self.dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((61, 15, 61, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(61, 29, 61, 33): 'self', (61, 35, 61, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((82, 15, 82, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(82, 29, 82, 33): 'self', (82, 35, 82, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((101, 15, 101, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(101, 29, 101, 33): 'self', (101, 35, 101, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((110, 18, 110, 44), 'numpy.logical_not', 'np.logical_not', ({(110, 33, 110, 43): 'self.walls'}, {}), '(self.walls)', True, 'import numpy as np\n'), ((125, 15, 125, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(125, 29, 125, 33): 'self', (125, 35, 125, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((140, 15, 140, 38), 'ciabatta.meta.make_repr_str', 'make_repr_str', ({(140, 29, 140, 33): 'self', (140, 35, 140, 37): 'fs'}, {}), '(self, fs)', False, 'from ciabatta.meta import make_repr_str\n'), ((150, 8, 150, 42), 'fealty.field_numerics.density_1d', 'field_numerics.density_1d', ({(150, 34, 150, 38): 'inds', (150, 40, 150, 41): 'f'}, {}), '(inds, f)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((164, 8, 164, 53), 'fealty.field_numerics.laplace_1d', 'field_numerics.laplace_1d', ({(164, 34, 164, 39): 'field', (164, 41, 164, 48): 'laplace', (164, 50, 164, 52): 'dx'}, {}), '(field, laplace, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((180, 8, 180, 57), 'fealty.field_numerics.grad_i_1d', 'field_numerics.grad_i_1d', ({(180, 33, 180, 38): 'field', (180, 40, 180, 44): 'inds', (180, 46, 180, 52): 'grad_i', (180, 54, 180, 56): 'dx'}, {}), '(field, inds, grad_i, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((194, 8, 194, 47), 'fealty.field_numerics.grad_1d', 'field_numerics.grad_1d', ({(194, 31, 194, 36): 'field', (194, 38, 194, 42): 'grad', (194, 44, 194, 46): 'dx'}, {}), '(field, grad, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((208, 8, 208, 45), 'fealty.field_numerics.div_1d', 'field_numerics.div_1d', ({(208, 30, 208, 35): 'field', (208, 37, 208, 40): 'div', (208, 42, 208, 44): 'dx'}, {}), '(field, div, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((223, 8, 223, 61), 'fealty.walled_field_numerics.grad_1d', 'walled_field_numerics.grad_1d', ({(223, 38, 223, 43): 'field', (223, 45, 223, 49): 'grad', (223, 51, 223, 53): 'dx', (223, 55, 223, 60): 'walls'}, {}), '(field, grad, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((240, 8, 240, 71), 'fealty.walled_field_numerics.grad_i_1d', 'walled_field_numerics.grad_i_1d', ({(240, 40, 240, 45): 'field', (240, 47, 240, 51): 'inds', (240, 53, 240, 59): 'grad_i', (240, 61, 240, 63): 'dx', (240, 65, 240, 70): 'walls'}, {}), '(field, inds, grad_i, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((255, 8, 255, 67), 'fealty.walled_field_numerics.laplace_1d', 'walled_field_numerics.laplace_1d', ({(255, 41, 255, 46): 'field', (255, 48, 255, 55): 'laplace', (255, 57, 255, 59): 'dx', (255, 61, 255, 66): 'walls'}, {}), '(field, laplace, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((68, 17, 68, 62), 'numpy.ones', 'np.ones', (), '', True, 'import numpy as np\n'), ((152, 8, 152, 42), 'fealty.field_numerics.density_2d', 'field_numerics.density_2d', ({(152, 34, 152, 38): 'inds', (152, 40, 152, 41): 'f'}, {}), '(inds, f)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((166, 8, 166, 53), 'fealty.field_numerics.laplace_2d', 'field_numerics.laplace_2d', ({(166, 34, 166, 39): 'field', (166, 41, 166, 48): 'laplace', (166, 50, 166, 52): 'dx'}, {}), '(field, laplace, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((182, 8, 182, 57), 'fealty.field_numerics.grad_i_2d', 'field_numerics.grad_i_2d', ({(182, 33, 182, 38): 'field', (182, 40, 182, 44): 'inds', (182, 46, 182, 52): 'grad_i', (182, 54, 182, 56): 'dx'}, {}), '(field, inds, grad_i, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((196, 8, 196, 47), 'fealty.field_numerics.grad_2d', 'field_numerics.grad_2d', ({(196, 31, 196, 36): 'field', (196, 38, 196, 42): 'grad', (196, 44, 196, 46): 'dx'}, {}), '(field, grad, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((210, 8, 210, 45), 'fealty.field_numerics.div_2d', 'field_numerics.div_2d', ({(210, 30, 210, 35): 'field', (210, 37, 210, 40): 'div', (210, 42, 210, 44): 'dx'}, {}), '(field, div, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((225, 8, 225, 61), 'fealty.walled_field_numerics.grad_2d', 'walled_field_numerics.grad_2d', ({(225, 38, 225, 43): 'field', (225, 45, 225, 49): 'grad', (225, 51, 225, 53): 'dx', (225, 55, 225, 60): 'walls'}, {}), '(field, grad, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((242, 8, 242, 71), 'fealty.walled_field_numerics.grad_i_2d', 'walled_field_numerics.grad_i_2d', ({(242, 40, 242, 45): 'field', (242, 47, 242, 51): 'inds', (242, 53, 242, 59): 'grad_i', (242, 61, 242, 63): 'dx', (242, 65, 242, 70): 'walls'}, {}), '(field, inds, grad_i, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((257, 8, 257, 67), 'fealty.walled_field_numerics.laplace_2d', 'walled_field_numerics.laplace_2d', ({(257, 41, 257, 46): 'field', (257, 48, 257, 55): 'laplace', (257, 57, 257, 59): 'dx', (257, 61, 257, 66): 'walls'}, {}), '(field, laplace, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((154, 8, 154, 42), 'fealty.field_numerics.density_3d', 'field_numerics.density_3d', ({(154, 34, 154, 38): 'inds', (154, 40, 154, 41): 'f'}, {}), '(inds, f)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((168, 8, 168, 53), 'fealty.field_numerics.laplace_3d', 'field_numerics.laplace_3d', ({(168, 34, 168, 39): 'field', (168, 41, 168, 48): 'laplace', (168, 50, 168, 52): 'dx'}, {}), '(field, laplace, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((184, 8, 184, 51), 'fealty.field_numerics.grad_i_3d', 'field_numerics.grad_i_3d', ({(184, 33, 184, 38): 'field', (184, 40, 184, 46): 'grad_i', (184, 48, 184, 50): 'dx'}, {}), '(field, grad_i, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((198, 8, 198, 47), 'fealty.field_numerics.grad_3d', 'field_numerics.grad_3d', ({(198, 31, 198, 36): 'field', (198, 38, 198, 42): 'grad', (198, 44, 198, 46): 'dx'}, {}), '(field, grad, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((212, 8, 212, 45), 'fealty.field_numerics.div_3d', 'field_numerics.div_3d', ({(212, 30, 212, 35): 'field', (212, 37, 212, 40): 'div', (212, 42, 212, 44): 'dx'}, {}), '(field, div, dx)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((227, 8, 227, 61), 'fealty.walled_field_numerics.grad_3d', 'walled_field_numerics.grad_3d', ({(227, 38, 227, 43): 'field', (227, 45, 227, 49): 'grad', (227, 51, 227, 53): 'dx', (227, 55, 227, 60): 'walls'}, {}), '(field, grad, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((244, 8, 244, 71), 'fealty.walled_field_numerics.grad_i_3d', 'walled_field_numerics.grad_i_3d', ({(244, 40, 244, 45): 'field', (244, 47, 244, 51): 'inds', (244, 53, 244, 59): 'grad_i', (244, 61, 244, 63): 'dx', (244, 65, 244, 70): 'walls'}, {}), '(field, inds, grad_i, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n'), ((259, 8, 259, 67), 'fealty.walled_field_numerics.laplace_3d', 'walled_field_numerics.laplace_3d', ({(259, 41, 259, 46): 'field', (259, 48, 259, 55): 'laplace', (259, 57, 259, 59): 'dx', (259, 61, 259, 66): 'walls'}, {}), '(field, laplace, dx, walls)', False, 'from fealty import lattice, field_numerics, walled_field_numerics\n')]