
import numpy as np


# -----------------------------------------------------------------------------
#  DERIVATIVES OF PARABOLOID
# -----------------------------------------------------------------------------

def compute_paraboloid_points(P, a, b, c, d, e):
    """Computes the points of the paraboloid x(u,v) = (u, v, z(u,v)) with
    z(u,v) = a*u^2 + b*v^2 + c*u*v + d*u + e*v.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - a, b, c, d, e : float
        The parameters of the paraboloid.

    Returns:
    - x : np.array (n, 3)
        The points x(P), where the i-th row contains the (x,y,z) coordinates
        of the point x(p_i)
    """
    n = a*P[:,0]**2 + \
        b*P[:,1]**2 + \
        c*P[:,0]*P[:,1] + \
        d*P[:,0] + \
        e*P[:,1]
        
    x = np.dstack((P[:,0],P[:,1],n))[0]
    return x


def compute_paraboloid_first_derivatives(P, a, b, c, d, e):
    """Computes the first derivatives of the paraboloid x(u,v) = (u, v, z(u,v))
    with z(u,v) = a*u^2 + b*v^2 + c*u*v + d*u + e*v.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - a, b, c, d, e : float
        The parameters of the paraboloid.

    Returns:
    - x_u : np.array (n, 3)
        The vectors x_u(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_u(p_i).
    - x_v : np.array (n, 3)
        The vectors x_v(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_v(p_i).
    """
    nu = 2*a*P[:,0] + \
        c*P[:,1] + \
        d
    nv = 2*b*P[:,1] + \
        c*P[:,0] + \
        e
    
    z = np.zeros(P.shape[0])
    o = np.ones(P.shape[0])

    x_u = np.dstack((o,z,nu))[0]
    x_v = np.dstack((z,o,nv))[0]
    return x_u, x_v


def compute_paraboloid_second_derivatives(P, a, b, c, d, e):
    """Computes the second derivatives of the paraboloid x(u,v) = (u, v, z(u,v))
    with z(u,v) = a*u^2 + b*v^2 + c*u*v + d*u + e*v.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - a, b, c, d, e : float
        The parameters of the paraboloid.

    Returns:
    - x_uu : np.array (n, 3)
        The vectors x_uu(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_uu(p_i).
    - x_uv : np.array (n, 3)
        The vectors x_uv(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_uv(p_i).
    - x_vv : np.array (n, 3)
        The vectors x_vv(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_vv(p_i).
    """
    nuu = np.repeat(2*a, P.shape[0])
    nuv = np.repeat( c,  P.shape[0])
    nvv = np.repeat(2*b, P.shape[0])
            
    z = np.zeros(P.shape[0])
        
    x_uu = np.dstack((z,z,nuu))[0]
    x_uv = np.dstack((z,z,nuv))[0]
    x_vv = np.dstack((z,z,nvv))[0]
    
    return x_uu, x_uv, x_vv


# -----------------------------------------------------------------------------
#  DERIVATIVES OF TORUS
# -----------------------------------------------------------------------------

def compute_torus_points(P, R, r):
    """Computes the second derivatives of a torus.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - R : float
        The radius of revolution.
    - r : float
        The radius of the cross section.
    Returns:
    - x : np.array (n, 3)
        The points x(P), where the i-th row contains the (x,y,z) coordinates
        of the point x(p_i)
    """
    cosu = np.cos(P[:,0])
    sinu = np.sin(P[:,0])
    cosv = np.cos(P[:,1])
    sinv = np.sin(P[:,1])
    fx = (r * cosu+ R) * cosv
    fy = (r * cosu+ R) * sinv
    fz = r * sinu
        
    x = np.dstack((fx, fy, fz))[0]
    return x


def compute_torus_first_derivatives(P, R, r):
    """Computes the second derivatives of a torus.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - R : float
        The radius of revolution.
    - r : float
        The radius of the cross section.
    Returns:
    - x_u : np.array (n, 3)
        The vectors x_u(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_u(p_i).
    - x_v : np.array (n, 3)
        The vectors x_v(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_v(p_i).
    """
    cosu = np.cos(P[:,0])
    sinu = np.sin(P[:,0])
    cosv = np.cos(P[:,1])
    sinv = np.sin(P[:,1])

    z = np.zeros(P.shape[0])

    fx_u = - r * sinu * cosv
    fy_u = - r * sinu * sinv
    fz_u = r * cosu

    fx_v = -(r * cosu+ R) * sinv
    fy_v = (r * cosu+ R) * cosv
    fz_v = z
        
    x_u = np.dstack((fx_u, fy_u, fz_u))[0]
    x_v = np.dstack((fx_v, fy_v, fz_v))[0]
    return x_u, x_v


def compute_torus_second_derivatives(P, R, r):
    """Computes the second derivatives of a torus.
    Try to vectorize this function.

    Parameters:
    - P : np.array (n, 2)
        Contains in i-th row the (u, v) coordinates of the i-th parameter point p_i.
    - R : float
        The radius of revolution.
    - r : float
        The radius of the cross section.

    Returns:
    - x_uu : np.array (n, 3)
        The vectors x_uu(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_uu(p_i).
    - x_uv : np.array (n, 3)
        The vectors x_uv(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_uv(p_i).
    - x_vv : np.array (n, 3)
        The vectors x_vv(P), where the i-th row contains the (x,y,z) coordinates
        of the vector x_vv(p_i).
    """
    cosu = np.cos(P[:,0])
    sinu = np.sin(P[:,0])
    cosv = np.cos(P[:,1])
    sinv = np.sin(P[:,1])

    z = np.zeros(P.shape[0])

    fx_uu = - r * cosu * cosv
    fy_uu = - r * cosu * sinv
    fz_uu = - r * sinu

    fx_uv = (r * sinu) * sinv
    fy_uv = -(r * sinu) * cosv
    fz_uv = z

    fx_vv = -(r * cosu+ R) * cosv
    fy_vv = -(r * cosu+ R) * sinv
    fz_vv = z
        
    x_uu = np.dstack((fx_uu, fy_uu, fz_uu))[0]
    x_uv = np.dstack((fx_uv, fy_uv, fz_uv))[0]
    x_vv = np.dstack((fx_vv, fy_vv, fz_vv))[0]
    return x_uu, x_uv, x_vv


# -----------------------------------------------------------------------------
#  SHAPE OPERATOR
# -----------------------------------------------------------------------------

def compute_first_fundamental_form(x_u, x_v):
    """Computes the first fundamental form I.
    Try to vectorize this function.

    Parameters:
    - x_u : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_u(p_i).
    - x_v : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_v(p_i).

    Returns:
    - I : np.array (n, 2, 2)
        The first fundamental forms.
        The (i, j, k) position contains the (j, k) element of the first
        fundamental form I(p_i).
    """
    e00 = np.einsum('lm,ml->m',x_u.T,x_u)
    e01 = np.einsum('lm,ml->m',x_u.T,x_v)
    e10 = np.einsum('lm,ml->m',x_v.T,x_u)
    e11 = np.einsum('lm,ml->m',x_v.T,x_v)
    I = np.dstack((e00,e01,e10,e11))[0].reshape((x_u.shape[0],2,2))
    return I


def compute_surface_normal(x_u, x_v):
    """Computes the surface normal n.
    Try to vectorize this function.

    Parameters:
    - x_u : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_u(p_i).
    - x_v : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_v(p_i).

    Returns:
    - n : np.array (n, 3)
        The surface normals.
        The i-th row contains the (x,y,z) coordinates of the vector n(p_i).
    """
    exp = np.cross(x_u, x_v)
    div = np.linalg.norm(exp, axis=1)
    n = exp/div[:,None]
    return n


def compute_second_fundamental_form(x_uu, x_uv, x_vv, n):
    """Computes the second fundamental form II.
    Try to vectorize this function.

    Parameters:
    - x_uu : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_uu(p_i).
    - x_uv : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_uv(p_i).
    - x_vv : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_vv(p_i).
    - n : np.array (n, 3)
        The surface normals.
        The i-th row contains the (x,y,z) coordinates of the vector n(p_i).

    Returns:
    - II : np.array (n, 2, 2)
        The second fundamental forms.
        The (i, j, k) position contains the (j, k) element of the second
        fundamental form II(p_i).
    """
    e00 = np.einsum('lm,ml->m',n.T, x_uu)
    e01 = np.einsum('lm,ml->m',n.T, x_uv)
    e10 = np.einsum('lm,ml->m',n.T, x_uv)
    e11 = np.einsum('lm,ml->m',n.T, x_vv)
    II = np.dstack((e00,e01,e10,e11))[0].reshape((n.shape[0],2,2))
    return II


def compute_shape_operator(I, II):
    """Computes the shape operator S.
    Try to vectorize this function.

    Parameters:
    - I : np.array (n, 2, 2)
        The first fundamental forms.
        The (i, j, k) position contains the (j, k) element of the first
        fundamental form I(p_i).
    - II : np.array (n, 2, 2)
        The second fundamental forms.
        The (i, j, k) position contains the (j, k) element of the second
        fundamental form II(p_i).

    Returns:
    - S : np.array (n, 2, 2)
        The shape operators.
        The (i, j, k) position contains the (j, k) element of the shape
        operator S(p_i).
    """
    Iinv = np.linalg.inv(I)
    S = Iinv @ II
    return S


# -----------------------------------------------------------------------------
#  PRINCIPAL CURVATURES
# -----------------------------------------------------------------------------

def compute_principal_curvatures(S, x_u, x_v):
    """Computes principal curvatures and corresponding principal directions.
    Try to vectorize this function.

    Parameters:
    - S : np.array (n, 2, 2)
        The shape operators.
        The (i, j, k) position contains the (j, k) element of the shape
        operator S(p_i).
    - x_u : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_u(p_i).
    - x_v : np.array (n, 3)
        The i-th row contains the (x,y,z) coordinates of the vector x_v(p_i).

    Returns:
    - k_1 : np.array (n)
        The min principal curvature. i-th element contains the curvature k_1(p_i).
    - k_2 : np.array (n)
        The max principal curvature. i-th element contains the curvature k_2(p_i).
    - e_1 : np.array (n, 3)
        The unitized principal curvature direction corresponding to k_1.
        The i-th row contains the (x,y,z) coordinates of e_1(p_i).
    - e_2 : np.array (n, 3)
        The unitized principal curvature direction corresponding to k_2.
        The i-th row contains the (x,y,z) coordinates of e_2(p_i).
    """
    # this section computes the ordered eigenvalues and eigenvectors of S where
    # k_1[i] = min eigenvalue at p_i, k_2[i] = max eigenvalue at p_i,
    # bar_e_1[i] = [u, v] components of the eigenvector of k_1,
    # bar_e_2[i] = [u, v] components of the eigenvector of k_2
    eig = np.linalg.eig(S)
    index = np.argsort(eig[0], axis=1)
    k_1 = eig[0][np.arange(len(S)), index[:, 0]]
    k_2 = eig[0][np.arange(len(S)), index[:, 1]]
    bar_e_1 = eig[1][np.arange(len(S)), :, index[:, 0]]
    bar_e_2 = eig[1][np.arange(len(S)), :, index[:, 1]]

    # Compute the normalized 3D vectors e_1, e_2
    J = np.concatenate((x_u,x_v), axis=1).reshape((S.shape[0],2,3))
    Je1 = np.einsum('lnm,ln->lm',J, bar_e_1)
    Je2 = np.einsum('lnm,ln->lm',J, bar_e_2)
    e_1 = Je1/np.linalg.norm(Je1, axis=1)[:,None]
    e_2 = Je2/np.linalg.norm(Je2, axis=1)[:,None]
    return k_1, k_2, e_1, e_2


# -----------------------------------------------------------------------------
#  ASYMPTOTIC DIRECTIONS
# -----------------------------------------------------------------------------

def compute_asymptotic_directions(k_1, k_2, e_1, e_2):
    """Computes principal curvatures and corresponding principal directions.
    Try to vectorize this function.

    Parameters:
    - k_1 : np.array (n)
        The min principal curvature. i-th element contains the curvature k_1(p_i).
    - k_2 : np.array (n)
        The max principal curvature. i-th element contains the curvature k_2(p_i).
    - e_1 : np.array (n, 3)
        The unitized principal curvature direction corresponding to k_1.
        The i-th row contains the (x,y,z) coordinates of e_1(p_i).
    - e_2 : np.array (n, 3)
        The unitized principal curvature direction corresponding to k_2.
        The i-th row contains the (x,y,z) coordinates of e_2(p_i).

    Returns:
    - a_1 : np.array (n, 3)
        The first unitized asymptotic direction. The i-th row contains the
        (x,y,z) coordinates of a_2(p_i) if it exists, (0, 0, 0) otherwise.
    - a_2 : np.array (n, 3)
        The second unitized asymptotic direction. The i-th row contains the
        (x,y,z) coordinates of a_2(p_i) if it exists, (0, 0, 0) otherwise.
    """
    a_1 = np.zeros((k_1.shape[0],3))
    a_2 = np.zeros((k_1.shape[0],3))
    k1k2 = k_1*k_2
    v1 = np.where(k1k2<=0)
    v2 = np.where(k1k2<0)

    f = np.zeros((k_1.shape[0]))
    ff = np.zeros((k_1.shape[0]))
    v3 = np.where((k_1-k_2)!=0)

    f[v3] = (k_1[v3]/((k_1 - k_2)[v3]))
    ff[v1] = f[v1]

    cosT = np.sqrt(1-f)
    sinT = np.sqrt(f)

    a_1[v1] =(e_1*cosT[:,None] + e_2*sinT[:,None])[v1]
    a_2[v2] =(e_1*cosT[:,None] - e_2*sinT[:,None])[v2]
    return a_1, a_2

