#!/usr/bin/env python
"""Deform the lithosphere with 1D or 2D flexure.

Landlab component that implements a 1D lithospheric flexure.

Examples
--------

Create a grid on which we will run the flexure calculations.

>>> from landlab import RasterModelGrid
>>> from landlab.components.flexure import Flexure1D
>>> grid = RasterModelGrid((3, 4), spacing=(1.e4, 1.e4))

Because `Flexure1D` is a one-dimensional component, it operates
*independently* on each row of grid nodes. By default, it will
calculate deflections on every row but you can also specify
which rows to operate on by passing an indexing array to the
`rows` keyword.

In this example, we'll just calculate deflections along the
middle row or nodes (that is, row 1)

>>> flex = Flexure1D(grid, rows=1)

In creating the component, a field (initialized with zeros) was
added to the grid that represents the current loading distribution.
If the grid already had this field, the component would use the
existing field. This can be accessed either through the *grid*
attribute in the same way as other landlab fields,

>>> flex.grid.at_node['lithosphere__increment_of_overlying_pressure']
array([ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.])

or through the `load_at_node` attribute of `Flexure1D`,

>>> flex.load_at_node
array([[ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.]])

Notice that `load_at_node` returns a reshaped view of the array
whereas the field returns a flattened array. Change values in this
array to add loads to the grid,

>>> flex.load_at_node[1, 2:4] = flex.gamma_mantle
>>> flex.run_one_step()

The output deflections can be retrieved either using landlab fields
as,

>>> flex.grid.at_node['lithosphere_surface__increment_of_elevation']
array([ 0.,  0.,  0.,  0.,  0.,  0.,  1.,  1.,  0.,  0.,  0.,  0.])

or through the `dz_at_node` attribute,

>>> flex.dz_at_node
array([[ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  1.,  1.],
       [ 0.,  0.,  0.,  0.]])
"""

import numpy as np

from landlab import Component
from .funcs import get_flexure_parameter
from ...utils.decorators import use_file_name_or_kwds
from .ext import subside_load_1d


class Flexure1D(Component):

    """Deform the lithosphere with 1D flexure.

    Landlab component that implements a 1D lithospheric flexure model.

    Parameters
    ----------
    grid : RasterModelGrid
        A grid.
    eet : float, optional
        Effective elastic thickness (m).
    youngs : float, optional
        Young's modulus.
    method : {'airy', 'flexure'}, optional
        Method to use to calculate deflections.
    rho_mantle : float, optional
        Density of the mantle (kg / m^3).
    rho_water : float, optional
        Density of the sea water (kg / m^3).
    gravity : float, optional
        Acceleration due to gravity (m / s^2).
    rows : int, optional
        Node rows that this component will operate on (default is to
        operate on *all* rows).

    Examples
    --------
    >>> from landlab import RasterModelGrid
    >>> from landlab.components.flexure import Flexure
    >>> grid = RasterModelGrid((5, 4), spacing=(1.e4, 1.e4))

    >>> flex = Flexure(grid)
    >>> flex.name
    'Flexure'
    >>> flex.input_var_names
    ('lithosphere__overlying_pressure_increment',)
    >>> flex.output_var_names
    ('lithosphere_surface__elevation_increment',)
    >>> sorted(flex.units) # doctest: +NORMALIZE_WHITESPACE
    [('lithosphere__overlying_pressure_increment', 'Pa'),
     ('lithosphere_surface__elevation_increment', 'm')]

    >>> flex.grid.number_of_node_rows
    5
    >>> flex.grid.number_of_node_columns
    4
    >>> flex.grid is grid
    True

    >>> np.all(grid.at_node['lithosphere_surface__elevation_increment'] == 0.)
    True

    >>> np.all(grid.at_node['lithosphere__overlying_pressure_increment'] == 0.)
    True
    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere_surface__elevation_increment'] == 0.)
    True

    >>> load = grid.at_node['lithosphere__overlying_pressure_increment']
    >>> load[4] = 1e9
    >>> dz = grid.at_node['lithosphere_surface__elevation_increment']
    >>> np.all(dz == 0.)
    True

    >>> flex.update()
    >>> np.all(grid.at_node['lithosphere_surface__elevation_increment'] == 0.)
    False
    """

    _name = "1D Flexure Equation"

    _input_var_names = ("lithosphere__increment_of_overlying_pressure",)

    _output_var_names = ("lithosphere_surface__increment_of_elevation",)

    _var_units = {
        "lithosphere__increment_of_overlying_pressure": "Pa",
        "lithosphere_surface__increment_of_elevation": "m",
    }

    _var_mapping = {
        "lithosphere__increment_of_overlying_pressure": "node",
        "lithosphere_surface__increment_of_elevation": "node",
    }

    _var_doc = {
        "lithosphere__increment_of_overlying_pressure": "Applied pressure to the lithosphere over a time step",
        "lithosphere_surface__increment_of_elevation": "The change in elevation of the top of the lithosphere (the land "
        "surface) in one timestep",
    }

    POISSON = .25

    @use_file_name_or_kwds
    def __init__(
        self,
        grid,
        eet=65e3,
        youngs=7e10,
        method="airy",
        rho_mantle=3300.,
        rho_water=1030.,
        gravity=9.80665,
        rows=None,
        **kwds
    ):
        """Initialize the flexure component.

        Parameters
        ----------
        grid : RasterModelGrid
            A grid.
        eet : float, optional
            Effective elastic thickness (m).
        youngs : float, optional (Pa)
            Young's modulus.
        method : {'airy', 'flexure'}, optional
            Method to use to calculate deflections.
        rho_mantle : float, optional
            Density of the mantle (kg / m^3).
        rho_water : float, optional
            Density of the sea water (kg / m^3).
        gravity : float, optional
            Acceleration due to gravity (m / s^2).
        rows : int, optional
            Node rows that this component will operate on (default is to
            operate on *all* rows).
        """
        if method not in ("airy", "flexure"):
            raise ValueError("{method}: method not understood".format(method=method))

        self._grid = grid

        self._method = method
        self.youngs = youngs
        self.rho_mantle = rho_mantle
        self.rho_water = rho_water
        self.gravity = gravity
        self.eet = eet

        super(Flexure1D, self).__init__(grid, **kwds)

        for name in self._input_var_names + self._output_var_names:
            if name not in self.grid.at_node:
                self.grid.add_zeros(name, units=self._var_units[name], at="node")

        self._rows = (rows,) or Ellipsis

        self._x_at_node = self.grid.x_of_node.reshape(self.grid.shape).copy()

    @property
    def eet(self):
        """Effective elastic thickness (m)."""
        return self._eet

    @eet.setter
    def eet(self, new_val):
        if new_val <= 0:
            raise ValueError("Effective elastic thickness must be positive.")
        try:
            del self._rigidity, self._alpha
        except AttributeError:
            pass
        self._eet = float(new_val)

    @property
    def youngs(self):
        """Young's modulus of lithosphere (Pa)."""
        return self._youngs

    @youngs.setter
    def youngs(self, new_val):
        if new_val <= 0:
            raise ValueError("Young's modulus must be positive.")
        try:
            del self._rigidity, self._alpha
        except AttributeError:
            pass
        self._youngs = float(new_val)

    @property
    def rho_water(self):
        """Density of water (kg/m^3)."""
        return self._rho_water

    @rho_water.setter
    def rho_water(self, new_val):
        if new_val <= 0:
            raise ValueError("Water density must be positive.")
        try:
            del self._gamma_mantle, self._alpha
        except AttributeError:
            pass
        self._rho_water = float(new_val)

    @property
    def rho_mantle(self):
        """Density of mantle (kg/m^3)."""
        return self._rho_mantle

    @rho_mantle.setter
    def rho_mantle(self, new_val):
        if new_val <= 0:
            raise ValueError("Mantle density must be positive.")
        try:
            del self._gamma_mantle, self._alpha
        except AttributeError:
            pass
        self._rho_mantle = float(new_val)

    @property
    def gravity(self):
        """Acceleration due to gravity (m/s^2)."""
        return self._gravity

    @gravity.setter
    def gravity(self, new_val):
        if new_val <= 0:
            raise ValueError("Acceleration due to gravity must be positive.")
        try:
            del self._gamma_mantle, self._alpha
        except AttributeError:
            pass
        self._gravity = float(new_val)

    @property
    def alpha(self):
        """Flexure parameter (m)."""
        try:
            self._alpha
        except AttributeError:
            self._alpha = np.power(4 * self.rigidity / self.gamma_mantle, .25)
        finally:
            return self._alpha

    @property
    def rigidity(self):
        """Flexural rigidity (N m)."""
        try:
            self._rigidity
        except AttributeError:
            self._rigidity = (
                self._eet ** 3. * self._youngs / (12. * (1. - self.POISSON ** 2.))
            )
        finally:
            return self._rigidity

    @property
    def gamma_mantle(self):
        """Specific density of mantle (N/m^3)."""
        try:
            self._gamma_mantle
        except AttributeError:
            self._gamma_mantle = (self._rho_mantle - self._rho_water) * self._gravity
        finally:
            return self._gamma_mantle

    @property
    def method(self):
        """Name of method used to calculate deflections."""
        return self._method

    @property
    def x_at_node(self):
        return self._x_at_node

    @property
    def load_at_node(self):
        return self.grid.at_node[
            "lithosphere__increment_of_overlying_pressure"
        ].reshape(self.grid.shape)

    @property
    def dz_at_node(self):
        return self.grid.at_node["lithosphere_surface__increment_of_elevation"].reshape(
            self.grid.shape
        )

    def update(self):
        """Update fields with current loading conditions."""
        load = self.load_at_node[self._rows]
        deflection = self.dz_at_node[self._rows]

        if self._method == "airy":
            deflection[:] = load / self.gamma_mantle
        else:
            Flexure1D.calc_flexure(
                self.x_at_node[0], load, self.alpha, self.rigidity, out=deflection
            )

        if not np.may_share_memory(deflection, self.dz_at_node):
            self.dz_at_node[self._rows] = deflection

    def run_one_step(self, dt=None):
        self.update()

    def subside_loads(self, loads, out=None):
        """Subside surface due to multiple loads.

        Parameters
        ----------
        loads : ndarray of float
            Loads applied to each grid node (Pa).
        out : ndarray of float, optional
            Buffer to place resulting deflection values (m).

        Returns
        -------
        ndarray of float
            Deflections caused by the loading.
        """
        if self._method == "airy":
            deflection[:] = load / self.gamma_mantle
        else:
            Flexure1D.calc_flexure(
                self.x_at_node[0], loads, self.alpha, self.rigidity, out=out
            )

        return out

    @staticmethod
    def calc_flexure(x, loads, alpha, rigidity, out=None):
        """Subside surface due to multiple loads.

        Parameters
        ----------
        x : ndarray of float
            Position of grid nodes (m).
        loads : ndarray of float
            Loads applied to each grid node (Pa).
        alpha : float
            Flexure parameter of the lithosphere (m).
        rigidity : float
            Flexural rigidity of the lithosphere (N m).
        out : ndarray of float, optional
            Buffer to place resulting deflection values (m).

        Returns
        -------
        ndarray of float
            Deflections caused by the loading.
        """
        if out is None:
            out = np.zeros_like(loads, dtype=np.float)

        loads = loads.reshape((-1, loads.shape[-1]))
        out = out[..., :].reshape((-1, out.shape[-1]))
        subside_load_1d(x, loads, alpha, rigidity, out[..., :])

        if out.base is not None:
            return out.base
        else:
            return out
