'''
This module contains functions for building two and three-dimensional weight
matrices for linear elasticity problems.
'''
from rbf.pde.fd import weight_matrix

def elastic2d_body_force(x, p, n, lamb=1.0, mu=1.0, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the body force at
    `x` in a two-dimensional (plane strain) homogeneous elastic medium

    Parameters
    ----------
    x : (N, 2) array
        Target points.

    p : (M, 2) array
        Observation points.

    n : int
        stencil size

    lamb, mu : float, optional
        Lame parameters

    **kwargs :
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    # x component of force resulting from displacement in the x direction.
    coeffs_xx = [lamb + 2*mu, mu]
    diffs_xx = [(2, 0), (0, 2)]
    # x component of force resulting from displacement in the y direction.
    coeffs_xy = [lamb + mu]
    diffs_xy = [(1, 1)]
    # y component of force resulting from displacement in the x direction.
    coeffs_yx = [lamb + mu]
    diffs_yx = [(1, 1)]
    # y component of force resulting from displacement in the y direction.
    coeffs_yy = [lamb + 2*mu, mu]
    diffs_yy =  [(0, 2), (2, 0)]
    # make the differentiation matrices that enforce the PDE on the interior
    # nodes.
    D_xx = weight_matrix(x, p, n, diffs_xx, coeffs=coeffs_xx, **kwargs)
    D_xy = weight_matrix(x, p, n, diffs_xy, coeffs=coeffs_xy, **kwargs)
    D_yx = weight_matrix(x, p, n, diffs_yx, coeffs=coeffs_yx, **kwargs)
    D_yy = weight_matrix(x, p, n, diffs_yy, coeffs=coeffs_yy, **kwargs)
    return {'xx':D_xx, 'xy':D_xy, 'yx':D_yx, 'yy':D_yy}


def elastic2d_surface_force(x, nrm, p, n, lamb=1.0, mu=1.0, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the surface
    traction force at `x` with normals `nrm` in a two-dimensional (plane
    strain) homogeneous elastic medium.

    Parameters
    ----------
    x: (N, 2) array
        target points which reside on a surface.

    nrm: (N, 2) array
        surface normal vectors at each point in `x`.

    p: (M, 2) array
        observation points.

    n : int
        stencil size

    lamb, mu: float
        Lame parameters

    **kwargs:
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    # x component of traction force resulting from x displacement
    coeffs_xx = [nrm[:, 0]*(lamb + 2*mu), nrm[:, 1]*mu]
    diffs_xx =  [(1, 0), (0, 1)]
    # x component of traction force resulting from y displacement
    coeffs_xy = [nrm[:, 0]*lamb, nrm[:, 1]*mu]
    diffs_xy =  [(0, 1), (1, 0)]
    # y component of traction force resulting from x displacement
    coeffs_yx = [nrm[:, 0]*mu, nrm[:, 1]*lamb]
    diffs_yx =  [(0, 1), (1, 0)]
    # y component of force resulting from displacement in the y direction
    coeffs_yy = [nrm[:, 0]*mu, nrm[:, 1]*(lamb + 2*mu)]
    diffs_yy =  [(1, 0), (0, 1)]
    # make the differentiation matrices that enforce the free surface boundary
    # conditions.
    D_xx = weight_matrix(x, p, n, diffs_xx, coeffs=coeffs_xx, **kwargs)
    D_xy = weight_matrix(x, p, n, diffs_xy, coeffs=coeffs_xy, **kwargs)
    D_yx = weight_matrix(x, p, n, diffs_yx, coeffs=coeffs_yx, **kwargs)
    D_yy = weight_matrix(x, p, n, diffs_yy, coeffs=coeffs_yy, **kwargs)
    return {'xx':D_xx, 'xy':D_xy, 'yx':D_yx, 'yy':D_yy}


def elastic2d_displacement(x, p, n, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the displacements
    at `x`.

    Parameters
    ----------
    x: (N, 2) array
        target points.

    p: (M, 2) array
        observation points.

    n : int
        stencil size

    **kwargs:
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    D_xx = weight_matrix(x, p, n, (0, 0), **kwargs)
    D_yy = weight_matrix(x, p, n, (0, 0), **kwargs)
    return {'xx':D_xx, 'yy':D_yy}


def elastic3d_body_force(x, p, n, lamb=1.0, mu=1.0, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the body force at
    `x` in a three-dimensional homogeneous elastic medium.

    Parameters
    ----------
    x: (N, 3) array
        target points.

    p: (M, 3) array
        observation points.

    n : int
        stencil size

    lamb, mu: float
        first Lame parameter

    **kwargs:
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    coeffs_xx = [lamb + 2*mu, mu, mu]
    diffs_xx =  [(2, 0, 0), (0, 2, 0), (0, 0, 2)]
    coeffs_xy = [lamb + mu]
    diffs_xy =  [(1, 1, 0)]
    coeffs_xz = [lamb + mu]
    diffs_xz =  [(1, 0, 1)]
    coeffs_yx = [lamb + mu]
    diffs_yx =  [(1, 1, 0)]
    coeffs_yy = [ mu, lamb + 2*mu, mu]
    diffs_yy =  [(2, 0, 0), (0, 2, 0), (0, 0, 2)]
    coeffs_yz = [lamb + mu]
    diffs_yz =  [(0, 1, 1)]
    coeffs_zx = [lamb + mu]
    diffs_zx =  [(1, 0, 1)]
    coeffs_zy = [lamb + mu]
    diffs_zy =  [(0, 1, 1)]
    coeffs_zz = [mu, mu, lamb + 2*mu]
    diffs_zz =  [(2, 0, 0), (0, 2, 0), (0, 0, 2)]
    D_xx = weight_matrix(x, p, n, diffs_xx, coeffs=coeffs_xx, **kwargs)
    D_xy = weight_matrix(x, p, n, diffs_xy, coeffs=coeffs_xy, **kwargs)
    D_xz = weight_matrix(x, p, n, diffs_xz, coeffs=coeffs_xz, **kwargs)
    D_yx = weight_matrix(x, p, n, diffs_yx, coeffs=coeffs_yx, **kwargs)
    D_yy = weight_matrix(x, p, n, diffs_yy, coeffs=coeffs_yy, **kwargs)
    D_yz = weight_matrix(x, p, n, diffs_yz, coeffs=coeffs_yz, **kwargs)
    D_zx = weight_matrix(x, p, n, diffs_zx, coeffs=coeffs_zx, **kwargs)
    D_zy = weight_matrix(x, p, n, diffs_zy, coeffs=coeffs_zy, **kwargs)
    D_zz = weight_matrix(x, p, n, diffs_zz, coeffs=coeffs_zz, **kwargs)
    return {'xx':D_xx, 'xy':D_xy, 'xz':D_xz,
            'yx':D_yx, 'yy':D_yy, 'yz':D_yz,
            'zx':D_zx, 'zy':D_zy, 'zz':D_zz}


def elastic3d_surface_force(x, nrm, p, n, lamb=1.0, mu=1.0, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the surface
    traction force at `x` with normals `nrm` in a three-dimensional homogeneous
    elastic medium.

    Parameters
    ----------
    x: (N, 3) array
        target points which reside on a surface.

    nrm: (N, 3) array
        surface normal vectors at each point in `x`.

    p: (M, 3) array
        observation points.

    n : int
        stencil size

    lamb, mu: float
        Lame parameters

    **kwargs:
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    coeffs_xx = [nrm[:, 0]*(lamb + 2*mu), nrm[:, 1]*mu, nrm[:, 2]*mu]
    diffs_xx =  [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
    coeffs_xy = [nrm[:, 0]*lamb, nrm[:, 1]*mu]
    diffs_xy =  [(0, 1, 0), (1, 0, 0)]
    coeffs_xz = [nrm[:, 0]*lamb, nrm[:, 2]*mu]
    diffs_xz =  [(0, 0, 1), (1, 0, 0)]
    coeffs_yx = [nrm[:, 0]*mu, nrm[:, 1]*lamb]
    diffs_yx =  [(0, 1, 0), (1, 0, 0)]
    coeffs_yy = [nrm[:, 0]*mu, nrm[:, 1]*(lamb + 2*mu), nrm[:, 2]*mu]
    diffs_yy =  [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
    coeffs_yz = [nrm[:, 1]*lamb, nrm[:, 2]*mu]
    diffs_yz =  [(0, 0, 1), (0, 1, 0)]
    coeffs_zx = [nrm[:, 0]*mu, nrm[:, 2]*lamb]
    diffs_zx =  [(0, 0, 1), (1, 0, 0)]
    coeffs_zy = [nrm[:, 1]*mu, nrm[:, 2]*lamb]
    diffs_zy =  [(0, 0, 1), (0, 1, 0)]
    coeffs_zz = [nrm[:, 0]*mu, nrm[:, 1]*mu, nrm[:, 2]*(lamb + 2*mu)]
    diffs_zz =  [(1, 0, 0), (0, 1, 0), (0, 0, 1)]
    D_xx = weight_matrix(x, p, n, diffs_xx, coeffs=coeffs_xx, **kwargs)
    D_xy = weight_matrix(x, p, n, diffs_xy, coeffs=coeffs_xy, **kwargs)
    D_xz = weight_matrix(x, p, n, diffs_xz, coeffs=coeffs_xz, **kwargs)
    D_yx = weight_matrix(x, p, n, diffs_yx, coeffs=coeffs_yx, **kwargs)
    D_yy = weight_matrix(x, p, n, diffs_yy, coeffs=coeffs_yy, **kwargs)
    D_yz = weight_matrix(x, p, n, diffs_yz, coeffs=coeffs_yz, **kwargs)
    D_zx = weight_matrix(x, p, n, diffs_zx, coeffs=coeffs_zx, **kwargs)
    D_zy = weight_matrix(x, p, n, diffs_zy, coeffs=coeffs_zy, **kwargs)
    D_zz = weight_matrix(x, p, n, diffs_zz, coeffs=coeffs_zz, **kwargs)
    return {'xx':D_xx, 'xy':D_xy, 'xz':D_xz,
            'yx':D_yx, 'yy':D_yy, 'yz':D_yz,
            'zx':D_zx, 'zy':D_zy, 'zz':D_zz}


def elastic3d_displacement(x, p, n, **kwargs):
    '''
    Returns weight matrices that map displacements at `p` to the displacements
    at `x`.

    Parameters
    ----------
    x: (N, 3) array
        target points.

    p: (M, 3) array
        observation points.

    n : int
        stencil size

    **kwargs:
        additional arguments passed to `weight_matrix`

    Returns
    -------
    dict
        keys are the components and the values are the corresponding weight
        matrices.

    '''
    D_xx = weight_matrix(x, p, n, (0, 0, 0), **kwargs)
    D_yy = weight_matrix(x, p, n, (0, 0, 0), **kwargs)
    D_zz = weight_matrix(x, p, n, (0, 0, 0), **kwargs)
    return {'xx':D_xx, 'yy':D_yy, 'zz':D_zz}
