#!/usr/bin/env python

import numpy as np
import scipy.special
from multiprocessing import Pool

_POISSON = .25

_N_PROCS = 4


def get_flexure_parameter(h, E, n_dim, gamma_mantle=33000.):
    """
    Calculate the flexure parameter based on some physical constants. *h* is
    the Effective elastic thickness of Earth's crust (m), *E* is Young's
    Modulus, and *n_dim* is the number of spatial dimensions for which the
    flexure parameter is used. The number of dimension must be either 1, or
    2.

    Examples
    --------
    >>> from __future__ import print_function
    >>> from landlab.components.flexure import get_flexure_parameter

    >>> eet = 65000.
    >>> youngs = 7e10
    >>> alpha = get_flexure_parameter(eet, youngs, 1)
    >>> print('%.3f' % round(alpha, 3))
    119965.926

    >>> alpha = get_flexure_parameter(eet, youngs, 2)
    >>> print('%.2f' % alpha)
    84828.72
    """
    D = E * pow(h, 3) / 12. / (1. - pow(_POISSON, 2))

    assert(n_dim == 1 or n_dim == 2)

    if n_dim == 2:
        alpha = pow(D / gamma_mantle, .25)
    else:
        alpha = pow(4. * D / gamma_mantle, .25)

    return alpha


def _calculate_distances(locs, coords):
    if isinstance(locs[0], (float, int)):
        return np.sqrt(pow(coords[0] - locs[0], 2) +
                       pow(coords[1] - locs[1], 2))
    else:
        r = pow(coords[0][:, np.newaxis] - locs[0], 2)
        r += pow(coords[1][:, np.newaxis] - locs[1], 2)
        return np.sqrt(r, out=r)


def _calculate_deflections(load, locs, coords, alpha, out=None,
                           gamma_mantle=33000.):
    c = - load / (2. * np.pi * gamma_mantle * pow(alpha, 2.))
    r = _calculate_distances(locs, coords) / alpha

    if isinstance(c, (float, int)):
        return np.multiply(scipy.special.kei(r), c, out=out)
    else:
        scipy.special.kei(r, out=r)
        np.multiply(r, c[np.newaxis, :], out=r)
        return np.sum(r, axis=1, out=out)


def subside_point_load(load, loc, coords, params=None, out=None):
    """Calculate deflection at points due a point load.

    Calculate deflections on a grid, defined by the points in the *coords*
    tuple, due to a point load of magnitude *load* applied at *loc*.

    *x* and *y* are the x and y coordinates of each node of the solution
    grid (in meters). The scalars *eet* and *youngs* define the crustal
    properties.

    Parameters
    ----------
    load : float
        Magnitude of the point load.
    loc : float or tuple
        Location of the load as either a scalar or as (*x*, *y*)
    coords : ndarray
        Array of points to calculate deflections at
    params : dict-like
        Physical parameters used for deflection calculation. Valid keys are
        - *eet*: Effective elastic thickness
        - *youngs*: Young's modulus
    out : ndarray, optional
        Array to put deflections into.

    Returns
    -------
    out : ndarray
        Array of deflections.

    Examples
    --------

    >>> from landlab.components.flexure import subside_point_load