# Modified from code presented in:
# Nikkhoo, M., Walter, T. R. (2015): Triangular dislocation: an analytical,
# artefact-free solution. - Geophysical Journal International, 201,
# 1117-1139. doi: 10.1093/gji/ggv035

# Original documentation:
# TDdispFS
# calculates displacements associated with a triangular dislocation in an
# elastic full-space.
#
# TD: Triangular Dislocation
# EFCS: Earth-Fixed Coordinate System
# TDCS: Triangular Dislocation Coordinate System
# ADCS: Angular Dislocation Coordinate System
#
# INPUTS
# X, Y and Z:
# Coordinates of calculation points in EFCS (East, North, Up). X, Y and Z
# must have the same size.
#
# P1,P2 and P3:
# Coordinates of TD vertices in EFCS.
#
# Ss, Ds and Ts:
# TD slip vector components (Strike-slip, Dip-slip, Tensile-slip).
#
# nu:
# Poisson's ratio.
#
# OUTPUTS
# ue, un and uv:
# Calculated displacement vector components in EFCS. ue, un and uv have
# the same unit as Ss, Ds and Ts in the inputs.
#
# Original documentation license:
# Copyright (c) 2014 Mehdi Nikkhoo
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.

import numpy as np

def build_tri_coordinate_system(tri):
    Vnorm = normalize(np.cross(tri[1] - tri[0], tri[2] - tri[0]))
    eY = np.array([0, 1, 0])
    eZ = np.array([0, 0, 1])
    Vstrike = np.cross(eZ, Vnorm)
    if np.linalg.norm(Vstrike) == 0:
        Vstrike = eY * Vnorm[2]
    Vstrike = normalize(Vstrike)
    Vdip = np.cross(Vnorm, Vstrike)
    return np.array([Vnorm, Vstrike, Vdip])

def normalize(v):
    return v / np.linalg.norm(v)

def trimodefinder(obs, tri):
    # trimodefinder calculates the normalized barycentric coordinates of
    # the points with respect to the TD vertices and specifies the appropriate
    # artefact-free configuration of the angular dislocations for the
    # calculations. The input matrices x, y and z share the same size and
    # correspond to the y, z and x coordinates in the TDCS, respectively. p1,
    # p2 and p3 are two-component matrices representing the y and z coordinates
    # of the TD vertices in the TDCS, respectively.
    # The components of the output (trimode) corresponding to each calculation
    # points, are 1 for the first configuration, -1 for the second
    # configuration and 0 for the calculation point that lie on the TD sides.

    a = ((tri[1,1]-tri[2,1])*(obs[0]-tri[2,0])+(tri[2,0]-tri[1,0])*(obs[1]-tri[2,1]))/ \
        ((tri[1,1]-tri[2,1])*(tri[0,0]-tri[2,0])+(tri[2,0]-tri[1,0])*(tri[0,1]-tri[2,1]))
    b = ((tri[2,1]-tri[0,1])*(obs[0]-tri[2,0])+(tri[0,0]-tri[2,0])*(obs[1]-tri[2,1]))/ \
        ((tri[1,1]-tri[2,1])*(tri[0,0]-tri[2,0])+(tri[2,0]-tri[1,0])*(tri[0,1]-tri[2,1]))
    c = 1-a-b

    result = 1
    if (a<=0 and b>c and c>a) or\
            (b<=0 and c>a and a>b) or\
            (c<=0 and a>b and b>c):
        result = -1
    if (a==0 and b>=0 and c>=0) or\
            (a>=0 and b==0 and c>=0) or\
            (a>=0 and b>=0 and c==0):
        result = 0
    if result == 0 and obs[2] != 0:
        result = 1
    return result


def TDSetupD(obs, alpha, slip_b, nu, TriVertex, SideVec):
    # TDSetupD transforms coordinates of the calculation points as well as
    # slip vector components from ADCS into TDCS. It then calculates the
    # displacements in ADCS and transforms them into TDCS.

    # Transformation matrix
    A = np.array([[SideVec[2], -SideVec[1]], [SideVec[1], SideVec[2]]])

    # Transform coordinates of the calculation points from TDCS into ADCS
    r1 = A.dot([obs[1]-TriVertex[1], obs[2]-TriVertex[2]])
    y1 = r1[0]
    z1 = r1[1]

    # Transform the in-plane slip vector components from TDCS into ADCS
    r2 = A.dot([slip_b[1], slip_b[2]])
    by1 = r2[0]
    bz1 = r2[1]

    # Calculate displacements associated with an angular dislocation in ADCS
    [u,v0,w0] = AngDisDisp(obs[0],y1,z1,-np.pi+alpha,slip_b[0],by1,bz1,nu)

    # Transform displacements from ADCS into TDCS
    r3 = A.T.dot([v0,w0])
    v = r3[0]
    w = r3[1]
    return u, v, w

def AngDisDisp(x, y, z, alpha, bx, by, bz, nu):
    # AngDisDisp calculates the "incomplete" displacements (without the
    # Burgers' function contribution) associated with an angular dislocation in
    # an elastic full-space.

    cosA = np.cos(alpha)
    sinA = np.sin(alpha)
    eta = y*cosA-z*sinA
    zeta = y*sinA+z*cosA
    r = np.sqrt(x ** 2 + y ** 2 + z ** 2)

    # Avoid complex results for the logarithmic terms
    if zeta > r:
        zeta = r
    if z > r:
        z = r

    ux = bx/8/np.pi/(1-nu)*(x*y/r/(r-z)-x*eta/r/(r-zeta))
    vx = bx/8/np.pi/(1-nu)*(eta*sinA/(r-zeta)-y*eta/r/(r-zeta)+\
        y**2/r/(r-z)+(1-2*nu)*(cosA*np.log(r-zeta)-np.log(r-z)))
    wx = bx/8/np.pi/(1-nu)*(eta*cosA/(r-zeta)-y/r-eta*z/r/(r-zeta)-\
        (1-2*nu)*sinA*np.log(r-zeta))

    uy = by/8/np.pi/(1-nu)*(x**2*cosA/r/(r-zeta)-x**2/r/(r-z)-\
        (1-2*nu)*(cosA*np.log(r-zeta)-np.log(r-z)))
    vy = by*x/8/np.pi/(1-nu)*(y*cosA/r/(r-zeta)-sinA*cosA/(r-zeta)-y/r/(r-z))
    wy = by*x/8/np.pi/(1-nu)*(z*cosA/r/(r-zeta)-cosA**2/(r-zeta)+1/r)

    uz = bz*sinA/8/np.pi/(1-nu)*((1-2*nu)*np.log(r-zeta)-x**2/r/(r-zeta))
    vz = bz*x*sinA/8/np.pi/(1-nu)*(sinA/(r-zeta)-y/r/(r-zeta))
    wz = bz*x*sinA/8/np.pi/(1-nu)*(cosA/(r-zeta)-z/r/(r-zeta))

    return ux+uy+uz, vx+vy+vz, wx+wy+wz

def TDdispFS(obs, tri, slip, nu):
    transform = build_tri_coordinate_system(tri)

    transformed_obs = transform.dot(obs - tri[1])
    transformed_tri = np.zeros((3,3))
    transformed_tri[0,:] = transform.dot(tri[0] - tri[1])
    transformed_tri[2,:] = transform.dot(tri[2] - tri[1])
    np.testing.assert_almost_equal(transformed_tri[1], [0,0,0])
    np.testing.assert_almost_equal(transformed_tri[0][0], 0)
    np.testing.assert_almost_equal(transformed_tri[2][0], 0)

    e12 = normalize(transformed_tri[1] - transformed_tri[0])
    e13 = normalize(transformed_tri[2] - transformed_tri[0])
    e23 = normalize(transformed_tri[2] - transformed_tri[1])

    A = np.arccos(e12.T.dot(e13))
    B = np.arccos(-e12.T.dot(e23))
    C = np.arccos(e23.T.dot(e13))

    mode = trimodefinder(np.array([transformed_obs[1],transformed_obs[2],transformed_obs[0]]),transformed_tri[:,1:])

    slip_b = np.array([slip[2], slip[0], slip[1]])
    if mode == 1:
        # Calculate first angular dislocation contribution
        u1Tp,v1Tp,w1Tp = TDSetupD(transformed_obs,A,slip_b,nu,transformed_tri[0], -e13)
        # Calculate second angular dislocation contribution
        u2Tp,v2Tp,w2Tp = TDSetupD(transformed_obs,B,slip_b,nu,transformed_tri[1], e12)
        # Calculate third angular dislocation contribution
        u3Tp,v3Tp,w3Tp = TDSetupD(transformed_obs,C,slip_b,nu,transformed_tri[2], e23)
        out = np.array([
            u1Tp+u2Tp+u3Tp,
            v1Tp+v2Tp+v3Tp,
            w1Tp+w2Tp+w3Tp
        ])
    elif mode == -1:
        # Calculate first angular dislocation contribution
        u1Tn,v1Tn,w1Tn = TDSetupD(transformed_obs,A,slip_b,nu,transformed_tri[0],e13)
        # Calculate second angular dislocation contribution
        u2Tn,v2Tn,w2Tn = TDSetupD(transformed_obs,B,slip_b,nu,transformed_tri[1],-e12)
        # Calculate third angular dislocation contribution
        u3Tn,v3Tn,w3Tn = TDSetupD(transformed_obs,C,slip_b,nu,transformed_tri[2],-e23)
        out = np.array([
            u1Tn+u2Tn+u3Tn,
            v1Tn+v2Tn+v3Tn,
            w1Tn+w2Tn+w3Tn
        ])
    else:
        out = np.array([np.nan,np.nan,np.nan])


    a = np.array([
        -transformed_obs[0],
        transformed_tri[0][1] - transformed_obs[1],
        transformed_tri[0][2] - transformed_obs[2]
    ])
    b = -transformed_obs
    c = np.array([
        -transformed_obs[0],
        transformed_tri[2][1] - transformed_obs[1],
        transformed_tri[2][2] - transformed_obs[2]
    ])
    na = np.sqrt(np.sum(a**2))
    nb = np.sqrt(np.sum(b**2))
    nc = np.sqrt(np.sum(c**2))

    FiN = (a[0]*(b[1]*c[2]-b[2]*c[1])- \
           a[1]*(b[0]*c[2]-b[2]*c[0])+ \
           a[2]*(b[0]*c[1]-b[1]*c[0]))
    FiD = (na*nb*nc+np.sum(a*b)*nc+np.sum(a*c)*nb+np.sum(b*c)*na)
    Fi = -2*np.arctan2(FiN,FiD)/4/np.pi

    # Calculate the complete displacement vector components in TDCS
    out += slip_b * Fi

    # Transform the complete displacement vector components from TDCS into EFCS
    return transform.T.dot(out)
