"""
Compute the coordinate transformation related variables for a mapping,

    .. math::

        \\Phi : \\Omega_{\\mathrm{ref}}\\to\\Omega

We get an instance of class ``CoordinateTransformation``. The
transformations are its methods.

⭕ To access the source code, click on the [source] button at the right
side or click on
:download:`[coordinate_transformation.py]</contents/LIBRARY/ptc/mathischeap_ptc/coordinate_transformation.py>`.
Dependence may exist. In case of error, check import and install
required packages or download required scripts. © mathischeap.com

"""

import numpy as np
from numpy import sin as sin
from numpy import cos as cos
from numpy import pi as pi


def Phi(xi, et, sg):  # define a mapping; for testing
    r, s, t = (xi + 1) / 2, (et + 1) / 2, (sg + 1) / 2
    x = r + 0.1 * sin(2 * pi * r) * sin(2 * pi * s) * sin(2 * pi * t)
    y = s + 0.1 * sin(2 * pi * r) * sin(2 * pi * s) * sin(2 * pi * t)
    z = t + 0.1 * sin(2 * pi * r) * sin(2 * pi * s) * sin(2 * pi * t)
    return x, y, z


def d_Phi(xi, et, sg):  # the Jacobian matrix of above mapping
    r, s, t = (xi + 1) / 2, (et + 1) / 2, (sg + 1) / 2
    x_xi = (1 + 0.2*pi*cos(2*pi * r) * sin(2*pi * s) * sin(2*pi * t))/2
    x_et = (    0.2*pi*sin(2*pi * r) * cos(2*pi * s) * sin(2*pi * t))/2
    x_sg = (    0.2*pi*sin(2*pi * r) * sin(2*pi * s) * cos(2*pi * t))/2
    y_xi = (    0.2*pi*cos(2*pi * r) * sin(2*pi * s) * sin(2*pi * t))/2
    y_et = (1 + 0.2*pi*sin(2*pi * r) * cos(2*pi * s) * sin(2*pi * t))/2
    y_sg = (    0.2*pi*sin(2*pi * r) * sin(2*pi * s) * cos(2*pi * t))/2
    z_xi = (    0.2*pi*cos(2*pi * r) * sin(2*pi * s) * sin(2*pi * t))/2
    z_et = (    0.2*pi*sin(2*pi * r) * cos(2*pi * s) * sin(2*pi * t))/2
    z_sg = (1 + 0.2*pi*sin(2*pi * r) * sin(2*pi * s) * cos(2*pi * t))/2
    return (x_xi, x_et, x_sg), (y_xi, y_et, y_sg), (z_xi, z_et, z_sg)



class CoordinateTransformation(object):
    """ The coordinate transformation class.

    :param Phi: The mapping :math:`\\Phi`. It should be a function which
        returns three components of the mapping, i.e.,

        .. math::

            \\Phi = (\\Phi_x, \\Phi_y, \\Phi_z).

    :type Phi: function
    :param d_Phi: The Jacobian matrix of :math:`\\Phi`,
        :math:`\\mathcal{J}`. It should be
        a function return the 9 components of the Jacobian matrix, i.e.,

        (
        :math:`\\dfrac{\\partial x}{\\partial\\xi}`,
        :math:`\\dfrac{\\partial x}{\\partial\\eta}`,
        :math:`\\dfrac{\\partial x}{\\partial\\varsigma}`),
        (
        :math:`\\dfrac{\\partial y}{\\partial\\xi}`,
        :math:`\\dfrac{\\partial y}{\\partial\\eta}`,
        :math:`\\dfrac{\\partial y}{\\partial\\varsigma}`),
        (
        :math:`\\dfrac{\\partial z}{\\partial\\xi}`,
        :math:`\\dfrac{\\partial z}{\\partial\\eta}`,
        :math:`\\dfrac{\\partial z}{\\partial\\varsigma}`).
    :type d_Phi: function

    :example:

    >>> ct = CoordinateTransformation(Phi, d_Phi) # generate a CT instance
    >>> xi = np.linspace(-0.9, 0.6, 3)
    >>> et = np.linspace(-0.7, 0.9, 3)
    >>> sg = np.linspace(-0.8, 0.7, 3)
    >>> xi, et, sg = np.meshgrid(xi, et, sg, indexing='ij')
    >>> x, y, z = ct.mapping(xi, et, sg)
    >>> x # doctest: +ELLIPSIS
    array([[[0.06469463, 0.05391086, 0.02977458],...
    >>> y # doctest: +ELLIPSIS
    array([[[0.16469463, 0.15391086, 0.12977458],...
    >>> z # doctest: +ELLIPSIS
    array([[[0.11469463, 0.47891086, 0.82977458],...
    >>> J = ct.Jacobian_matrix(xi, et, sg)
    >>> J[0][0] # J_{1,1} # doctest: +ELLIPSIS
    array([[[0.64207986, 0.53781345, 0.30444385],...
    >>> J[1][1] # J_{2,2} # doctest: +ELLIPSIS
    array([[[0.53354051, 0.50892654, 0.45383545],...
    >>> detJ = ct.Jacobian(xi, et, sg)
    >>> detJ # doctest: +ELLIPSIS
    array([[[0.1847901 , 0.11729178, 0.07611096],...
    >>> iJ = ct.inverse_Jacobian_matrix(xi, et, sg)
    >>> iJ[2][2] # J^{-1}_{3,3} # doctest: +ELLIPSIS
    array([[[1.82807508, 2.33068327, 1.69672867],...
    >>> iJ[1][2] # J^{-1}_{2,3} # doctest: +ELLIPSIS
    array([[[-0.17192492,  0.33068327, -0.30327133],...
    >>> G = ct.metric_matrix(xi, et, sg)
    >>> G[0][1] # g_{1,2} # doctest: +ELLIPSIS
    array([[[ 0.10210648,  0.02438263, -0.09377707],...
    >>> G[1][2] # g_{2,3} # doctest: +ELLIPSIS
    array([[[ 0.05493377, -0.03640053, -0.0063935 ],...
    >>> g = ct.metric(xi, et, sg)
    >>> g # doctest: +ELLIPSIS
    array([[[0.03414738, 0.01375736, 0.00579288],...
    >>> iG = ct.inverse_metric_matrix(xi, et, sg)
    >>> iG[1][2] # g^{2,3} # doctest: +ELLIPSIS
    array([[[-0.33977063,  0.72204401,  1.83434454],...

    """
    def __init__(self, Phi, d_Phi):
        self.Phi = Phi
        self.d_Phi = d_Phi

    def mapping(self, xi, et, sg):
        """A wrap of the input mapping, ``Phi``, i.e., :math:`\\Phi`.

        :param xi: :math:`\\xi`.
        :param et: :math:`\\eta`.
        :param sg: :math:`\\varsigma`.
        :return: a tuple :math:`(x,y,z) = \\Phi(\\xi,\\eta,\\varsigma)`.
        """
        return self.Phi(xi, et, sg)

    def Jacobian_matrix(self, xi, et, sg):
        """A wrap of the input Jacobian matrix, ``d_Phi``, i.e.,
        :math:`\\mathcal{J}`.

        :return: Return the Jacobian matrix :math:`\\mathcal{J}`:

            .. math::
                \\mathcal{J} =
                \\begin{bmatrix}
                \\dfrac{\\partial x}{\\partial \\xi} &
                \\dfrac{\\partial x}{\\partial \\eta} &
                \\dfrac{\\partial x}{\\partial \\varsigma} \\\\
                \\dfrac{\\partial y}{\\partial \\xi} &
                \\dfrac{\\partial y}{\\partial \\eta} &
                \\dfrac{\\partial y}{\\partial \\varsigma} \\\\
                \\dfrac{\\partial z}{\\partial \\xi} &
                \\dfrac{\\partial z}{\\partial \\eta} &
                \\dfrac{\\partial z}{\\partial \\varsigma}
                \\end{bmatrix}\\,.
        """
        return self.d_Phi(xi, et, sg)

    def Jacobian(self, xi, et, sg):
        """Compute the Jacobian :math:`\\mathrm{det}(\\mathcal{J})` of
        the mapping.

        :return: The Jacobian :math:`\\mathrm{det}(\\mathcal{J})`.
        """
        J = self.Jacobian_matrix(xi, et, sg)
        return + J[0][0]*J[1][1]*J[2][2] + J[0][1]*J[1][2]*J[2][0] \
               + J[0][2]*J[1][0]*J[2][1] - J[0][0]*J[1][2]*J[2][1] \
               - J[0][1]*J[1][0]*J[2][2] - J[0][2]*J[1][1]*J[2][0]

    def metric_matrix(self, xi, et, sg):
        """Compute the metric matrix :math:`\\mathcal{G}`.

        :return: Return the metric matrix :math:`\\mathcal{G}`:

            .. math::
                \\mathcal{G} =
                \\begin{bmatrix}
                g_{1,1} &
                g_{1,2} &
                g_{1,3} \\\\
                g_{2,1} &
                g_{2,2} &
                g_{2,3} \\\\
                g_{3,1} &
                g_{3,2} &
                g_{3,3}
                \\end{bmatrix}\\,.

        """
        J = self.Jacobian_matrix(xi, et, sg)
        G = [[None for _ in range(3)] for __ in range(3)]
        for i in range(3):
            for j in range(i, 3):
                G[i][j] = J[0][i] * J[0][j]
                for l in range(1, 3):
                    G[i][j] += J[l][i] * J[l][j]
                if i != j:
                    G[j][i] = G[i][j]
        return G

    def metric(self, xi, et, sg):
        """The metric
        :math:`g:= \\mathrm{det}(\\mathcal{G}):=(\mathrm{det}(\\mathcal{J}))^2`.
        """
        detJ = self.Jacobian(xi, et, sg)
        return detJ ** 2

    def inverse_Jacobian_matrix(self, xi, et, sg):
        """Compute the Jacobian matrix of the inverse mapping
        :math:`\\Phi^{-1}`,

        .. math::

            (\\xi,\\eta,\\varsigma) = \\Phi^{-1}(x,y,z).

        :return: Return the inverse Jacobian matrix :
            :math:`\\mathcal{J}^{-1}`:

                .. math::
                    \\mathcal{J}^{-1} =
                    \\begin{bmatrix}
                    \\dfrac{\\partial \\xi}{\\partial x} &
                    \\dfrac{\\partial \\xi}{\\partial y} &
                    \\dfrac{\\partial \\xi}{\\partial z} \\\\
                    \\dfrac{\\partial \\eta}{\\partial x} &
                    \\dfrac{\\partial \\eta}{\\partial y} &
                    \\dfrac{\\partial \\eta}{\\partial z} \\\\
                    \\dfrac{\\partial \\varsigma}{\\partial x} &
                    \\dfrac{\\partial \\varsigma}{\\partial y} &
                    \\dfrac{\\partial \\varsigma}{\\partial z}
                    \\end{bmatrix}\\,.

        """
        J = self.Jacobian_matrix(xi, et, sg)
        Jacobian = + J[0][0]*J[1][1]*J[2][2] + J[0][1]*J[1][2]*J[2][0] \
                   + J[0][2]*J[1][0]*J[2][1] - J[0][0]*J[1][2]*J[2][1] \
                   - J[0][1]*J[1][0]*J[2][2] - J[0][2]*J[1][1]*J[2][0]
        reciprocalJacobian = 1 / Jacobian
        del Jacobian
        iJ00 = reciprocalJacobian * (J[1][1] * J[2][2] - J[1][2] * J[2][1])
        iJ01 = reciprocalJacobian * (J[2][1] * J[0][2] - J[2][2] * J[0][1])
        iJ02 = reciprocalJacobian * (J[0][1] * J[1][2] - J[0][2] * J[1][1])
        iJ10 = reciprocalJacobian * (J[1][2] * J[2][0] - J[1][0] * J[2][2])
        iJ11 = reciprocalJacobian * (J[2][2] * J[0][0] - J[2][0] * J[0][2])
        iJ12 = reciprocalJacobian * (J[0][2] * J[1][0] - J[0][0] * J[1][2])
        iJ20 = reciprocalJacobian * (J[1][0] * J[2][1] - J[1][1] * J[2][0])
        iJ21 = reciprocalJacobian * (J[2][0] * J[0][1] - J[2][1] * J[0][0])
        iJ22 = reciprocalJacobian * (J[0][0] * J[1][1] - J[0][1] * J[1][0])
        return [[iJ00, iJ01, iJ02],
                [iJ10, iJ11, iJ12],
                [iJ20, iJ21, iJ22]]

    def inverse_Jacobian(self, xi, et, sg):
        """Compute the inverse Jacobian, the determinant of the inverse
        Jacobian matrix, :math:`\\mathrm{det}(\\mathcal{J}^{-1})`."""
        iJ = self.inverse_Jacobian_matrix(xi, et, sg)
        # noinspection PyUnresolvedReferences
        return + iJ[0][0]*iJ[1][1]*iJ[2][2] + iJ[0][1]*iJ[1][2]*iJ[2][0]\
               + iJ[0][2]*iJ[1][0]*iJ[2][1] - iJ[0][0]*iJ[1][2]*iJ[2][1]\
               - iJ[0][1]*iJ[1][0]*iJ[2][2] - iJ[0][2]*iJ[1][1]*iJ[2][0]

    def inverse_metric_matrix(self, xi, et, sg):
        """Compute the inverse metric matrix,
        :math:`\\mathcal{G}^{-1}`.

        :return: Return the inverse metric matrix :
            :math:`\\mathcal{G}^{-1}`:

                .. math::
                    \\mathcal{G}^{-1} =
                    \\begin{bmatrix}
                    g^{1,1} &
                    g^{1,2} &
                    g^{1,3} \\\\
                    g^{2,1} &
                    g^{2,2} &
                    g^{2,3} \\\\
                    g^{3,1} &
                    g^{3,2} &
                    g^{3,3}
                    \\end{bmatrix}\\,.
        """
        iJ = self.inverse_Jacobian_matrix(xi, et, sg)
        iG = [[None for _ in range(3)] for __ in range(3)]
        for i in range(3):
            for j in range(i, 3):
                # noinspection PyTypeChecker
                iG[i][j] = iJ[i][0] * iJ[j][0]
                for l in range(1, 3):
                    iG[i][j] += iJ[i][l] * iJ[j][l]
                if i != j:
                    iG[j][i] = iG[i][j]
        return iG







if __name__ == '__main__':
    import doctest
    doctest.testmod()
