# Add the igl library to the modules search path
import sys, os
sys.path.insert(0, os.getcwd() + "/../")

import pyigl as igl
from math import atan2,pi,cos,sin

# Mesh
V = igl.eigen.MatrixXd()
F = igl.eigen.MatrixXi()

# Constrained faces id
b = igl.eigen.MatrixXi()

# Constrained faces representative vector
bc = igl.eigen.MatrixXd()

# Degree of the N-RoSy field
N = 4;

# Converts a representative vector per face in the full set of vectors that describe
# an N-RoSy field
def representative_to_nrosy(V, F, R, N, Y):
    B1 = igl.eigen.MatrixXd()
    B2 = igl.eigen.MatrixXd()
    B3 = igl.eigen.MatrixXd()

    igl.local_basis(V,F,B1,B2,B3)

    Y.resize(F.rows()*N, 3)

    for i in range (0,F.rows()):
        x = R.row(i) * B1.row(i).transpose()
        y = R.row(i) * B2.row(i).transpose()
        angle = atan2(y[0],x[0])

        for j in range(0,N):
            anglej = angle + 2*pi*j/float(N)
            xj = cos(anglej)
            yj = sin(anglej)
            Y.setRow(i*N+j, xj * B1.row(i) + yj * B2.row(i))

# Plots the mesh with an N-RoSy field and its singularities on top
# The constrained faces (b) are colored in red.
def plot_mesh_nrosy(viewer, V, F, N, PD1, S, b):
    # Clear the mesh
    viewer.data.clear()
    viewer.data.set_mesh(V,F)

    # Expand the representative vectors in the full vector set and plot them as lines
    avg = igl.avg_edge_length(V, F)
    Y = igl.eigen.MatrixXd()
    representative_to_nrosy(V, F, PD1, N, Y)

    B = igl.eigen.MatrixXd()
    igl.barycenter(V,F,B)

    Be = igl.eigen.MatrixXd(B.rows()*N,3)
    for i in range(0,B.rows()):
        for j in range(0,N):
            Be.setRow(i*N+j,B.row(i))

    viewer.data.add_edges(Be,Be+Y*(avg/2),igl.eigen.MatrixXd([[0,0,1]]))

    # Plot the singularities as colored dots (red for negative, blue for positive)
    for i in range(0,S.size()):
        if S[i] < -0.001:
            viewer.data.add_points(V.row(i),igl.eigen.MatrixXd([[1,0,0]]))
        elif S[i] > 0.001:
            viewer.data.add_points(V.row(i),igl.eigen.MatrixXd([[0,1,0]]));

    # Highlight in red the constrained faces
    C = igl.eigen.MatrixXd.Constant(F.rows(),3,1)
    for i in range(0,b.size()):
        C.setRow(b[i], igl.eigen.MatrixXd([[1, 0, 0]]))
    viewer.data.set_colors(C)

# It allows to change the degree of the field when a number is pressed
def key_down(viewer, key, modifier):
    global N
    if key >= ord('1') and key <= ord('9'):
        N = key - ord('0')

    R = igl.eigen.MatrixXd()
    S = igl.eigen.MatrixXd()

    igl.comiso.nrosy(V,F,b,bc,igl.eigen.MatrixXi(),igl.eigen.MatrixXd(),igl.eigen.MatrixXd(),N,0.5,R,S)
    plot_mesh_nrosy(viewer,V,F,N,R,S,b)

    return False

# Load a mesh in OFF format
igl.readOFF("../../tutorial/shared/bumpy.off", V, F);

# Threshold faces with high anisotropy
b  = igl.eigen.MatrixXi([[0]])
bc = igl.eigen.MatrixXd([[1,1,1]])

viewer = igl.viewer.Viewer()

# Interpolate the field and plot
key_down(viewer, ord('4'), 0)

# Plot the mesh
viewer.data.set_mesh(V, F)
viewer.callback_key_down = key_down

# Disable wireframe
viewer.core.show_lines = False

# Launch the viewer
viewer.launch()