import igl
igl.__file__

#https://github.com/libigl/libigl-tutorial-data

import numpy as np
import igl
import meshplot as mp
from scipy.spatial.transform import Rotation
import ipywidgets as iw
import time
from math import exp
import quaternion
import pickle

from joblib import Parallel, delayed
import contextlib
import joblib
from tqdm import tqdm

@contextlib.contextmanager
def tqdm_joblib(tqdm_object):
    """Context manager to patch joblib to report into tqdm progress bar given as argument"""
    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
        def __call__(self, *args, **kwargs):
            tqdm_object.update(n=self.batch_size)
            return super().__call__(*args, **kwargs)

    old_batch_callback = joblib.parallel.BatchCompletionCallBack
    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
    try:
        yield tqdm_object
    finally:
        joblib.parallel.BatchCompletionCallBack = old_batch_callback
        tqdm_object.close()

V, F = igl.read_triangle_mesh('data/arm.obj')
C,BE,_,_,_,_ = igl.read_tgf('data/arm.tgf')
W = igl.read_dmat('data/arm-weights.dmat')
# labels = np.load('data/hand.label.npy').astype(int)
# v -= v.min(axis=0)
# v /= v.max()
# mp.plot(V,F)

# V, F = igl.read_triangle_mesh('data/untitled.obj')
# mp.plot(V,F)
F

# TV = np.concatenate((V, C))
# TE = BE + V.shape[0]

# labels = np.concatenate((np.zeros(V.shape[0]), np.array(range(1,C.shape[0]+1)))).astype(int)
labels = np.array(range(BE.shape[0])).astype(int)

def circle_sum(q1, q2):
    return q1 + q2 if np.dot(quaternion.as_float_array(q1), quaternion.as_float_array(q2)) >= 0 else q1 - q2

def apply_weighted_rotation(vertices, wgts, quat, CoR, translation): #vertices, weights, rot (as quat), CoR, translation array 
    vnew = vertices.copy()
    for i in range(vertices.shape[0]):
        vi = vertices[i]
        wi = wgts[i]
        
        q = quaternion.as_quat_array([0,0,0,0])
        for j in range(BE.shape[0]): # #BE
            qj = quat[j]
            wij = wi[j]
            q = circle_sum(q, wij*qj)

        # wqi  = np.sum(qi, axis=0)
        q /= np.linalg.norm(quaternion.as_float_array(q))
        
        R = quaternion.as_rotation_matrix(q)
        
        Rtilda = np.zeros((3,3))
        ttilda = np.zeros((3,1))
        
        for j in range(BE.shape[0]):
            Rj = quaternion.as_rotation_matrix(quat[j]) #3 x 3
            tj = translation[j] # 1 x 3
            wij = wi[j]
            Rtilda += wij*Rj
            ttilda += wij*tj.reshape((3,1))
            
        t = Rtilda @ CoR[i].reshape((3,1)) + ttilda - R @ CoR[i].reshape((3,1))
        
        vnew[i] = (R @ vi.reshape((3,1)) + t).reshape(3)

    return vnew

def similarity(Wp, Wv, sigma=0.1):
    #BE shape
    Wp = Wp if len(Wp.shape) == 1 else Wp.reshape(-1)
    Wv = Wv if len(Wv.shape) == 1 else Wv.reshape(-1)
    
    tot = 0
    for j in range(Wp.shape[0]):
        for k in range(j+1, Wv.shape[0]):
            tot += Wp[j]*Wp[k]*Wv[j]*Wv[k]*exp( -(Wp[j]*Wv[k] - Wp[k]*Wv[j])**2 / sigma**2)
            
    return tot

def CoR(i, weights, vertices, faces, omega=0.1, pb = None):
    
    num = np.zeros([1, 3])
    denom = np.zeros([1,3])
    for t in range(faces.shape[0]):
        # cmp = [i for i in range(3) if np.linalg.norm(weights[i] - weights[faces[t,i]]) < omega]
        cmp = [0, 1,2]
        if not cmp:
            continue
        s = similarity(weights[i], np.mean([weights[faces[t,c]] for c in cmp], axis = 0))
        v = np.mean([vertices[faces[t,c]] for c in cmp], axis = 0)
        a = igl.doublearea(vertices, faces[[t]]) / 2 

        num += s * v * a
        denom += s * a
    
    pi = num / denom
    
    if pb:
        pb.value = i

    return pi

# P = np.zeros(V.shape)
# with tqdm_joblib(tqdm(desc="My calculation", total=P.shape[0])) as progress_bar:
#     P = np.array(Parallel(n_jobs=10)(delayed(CoR)(i, W, V, F) for i in range(P.shape[0])))

# pickle.dump(P, open("data/CoR.p", "wb"))

P = pickle.load(open('data/CoR.p', 'rb'))
mp.plot(P)

OCoR = V.copy()

pos_f_saver = np.zeros((labels.max()+1, 7))

def pos_f(s,x,y,z, w, α, β, γ, plt):
    global OCoR

    clicked_edges = BE[np.isin(labels, s)]
    slices = []
    for e in clicked_edges:
        for v in e:
            if v not in slices:
                slices.append(v)
                
 
        
    # r = Rotation.from_euler('xyz', [α, β, γ], degrees=True)
    r = Rotation.from_quat([α, β, γ, w])
    
    v_slice = C[slices] + np.array([[x,y,z]])
    
    C1 = C.copy()
    C1[slices] = r.apply(v_slice)
    for si in s:
        pos_f_saver[si] = [x,y,z,w,α,β,γ]

    
    q = np.zeros((BE.shape[0], 4))
    t = np.zeros((BE.shape[0], 3))
    for i in range(BE.shape[0]):
        # res =  Qt0[i].inverse() * Qt1[i] 
        # diff[i,:] = [res.w, res.x, res.y, res.z]
        xi,yi,zi,wi,αi,βi,γi = pos_f_saver[i]
        q[i,:] = [wi,αi,βi,γi]
        t[i,:] = [xi, yi, zi]
    
    # q /= np.linalg.norm(q, axis=1)
    Q = quaternion.as_quat_array(q)
    
    OCoR = V.copy()
    OCoR = apply_weighted_rotation(V, W, Q, P, t)
    

    v_deformed = pos_f.deformer(OCoR)
    plt.update_object(vertices = v_deformed)
    # p.remove_object(max(p._Viewer__objects.keys()))
    # p.add_edges(C1+np.repeat([[0,.25,0]], C1.shape[0], axis=0), BE, shading={"line_color": "green"})
  
pos_f.deformer = lambda x:x

def widgets_wrapper():
    # segment_widget = iw.Dropdown(options=np.arange(labels.max()) + 1)
    segment_widget = iw.SelectMultiple(options=np.arange(labels.max()+1))
    translate_widget = {i:iw.FloatSlider(min=-1, max=1, value=0) 
                        for i in 'xyz'}
    # rotate_widget = {a:iw.FloatSlider(min=-90, max=90, value=0, step=1) 
    #                  for a in 'αβγ'}
    real_widget = {a:iw.FloatSlider(min=-2, max=2, value=1, step=.25) 
                     for a in 'w'}
    imag_widget = {a:iw.FloatSlider(min=-1, max=1, value=0, step=.1) 
                     for a in 'αβγ'}

    def update_seg(*args):
        (translate_widget['x'].value,translate_widget['y'].value,
        translate_widget['z'].value, real_widget['w'].value,
        imag_widget['α'].value,imag_widget['β'].value,
        imag_widget['γ'].value) = pos_f_saver[segment_widget.value]
    segment_widget.observe(update_seg, 'value')
    widgets_dict = dict(s=segment_widget)
    widgets_dict.update(translate_widget)
    widgets_dict.update(real_widget)
    widgets_dict.update(imag_widget)
    return widgets_dict

def position_deformer(target_pos):
    '''Fill in this function to change positions'''
    return target_pos
''' (Optional) Register this function to perform interactive deformation
pos_f.deformer = position_deformer
'''

## Widget UI

p = mp.plot(V, F)
# p.add_edges(C,BE, shading={"line_color": "green"});
iw.interact(pos_f,
            **widgets_wrapper(), plt=iw.fixed(p))

# deform(hand_full_rotation | {'x': .2}, p1)

def reset(plot):
    hand_full_rotation = {'s' : [0,1,2,3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : 1, 'α' : 0, 'β' : 0, 'γ' : 0}
    pos_f(**hand_full_rotation, plt=plot)

p1 = mp.plot(V,F)
hand_full_rotation = {'s' : [0,1,2,3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : 1, 'α' : 0, 'β' : 0, 'γ' : 0, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5


p1.add_mesh(v_deformed + shft, F)
p1.add_mesh(U+2*shft, F)
reset(p1)

p1 = mp.plot(V,F)
reset(p1)
hand_full_rotation = {'s' : [0,1,2,3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : 1, 'α' : 0, 'β' : 0, 'γ' : 0, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5


p1.add_mesh(v_deformed + shft, F)
p1.add_mesh(U+2*shft, F)

p1 = mp.plot(V,F)
reset(p1)
hand_full_rotation = {'s' : [0,1,2,3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : 0, 'α' : 0, 'β' : 0, 'γ' : -1, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5


# p1.add_mesh(v_deformed + shft, F)
p1.add_mesh(U+2*shft, F)

p1._Viewer__s["scale"] *= 2

# p1._Viewer__objects[0]["mesh"].add

import pythreejs as p3s
def add_text(viewer, text):
    sh = viewer._Viewer__get_shading({})
    tt = p3s.TextTexture(string=text, color="white", size=100, fontFace="cambria")
    sm = p3s.SpriteMaterial(map=tt, fog=False, transparent=False)
    
    obj = viewer._Viewer__objects[viewer._Viewer__cnt - 1]
    mesh = obj["mesh"]
    
    loc = (obj.get("min")[0], obj.get("max")[1], obj.get("max")[2])
    # loc[2] += .1
    print(text, viewer._Viewer__cnt, loc)
    
    s = p3s.Sprite(material=sm, scale = (.2,.2,.2), position=loc, center = (.5,.5))
    viewer.text = s
    viewer._scene.add(viewer.text)
    
#     sobj = {"mesh" : s}
    
#     viewer._Viewer__add_object(sobj, mesh)
    
    # p1._Viewer__update_view()

vtmp = np.zeros((1,3))
ftmp = np.zeros((1,3))

p1._Viewer__cnt

p1 = mp.plot(vtmp, ftmp)
reset(p1)

hand_full_rotation = {'s' : [2,3], 'x': -.1, 'y' : 0.2, 'z' : 0.0, 'w' : .5, 'α' : -1, 'β' : -.5, 'γ' : 0, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)
p1.add_mesh(v_deformed, F)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5

add_text(p1, "CoR")

# p1.add_mesh(v_deformed + shft, F)
p1.add_mesh(U+2*shft, F)

#adding in
add_text(p1, "DQS")

p1 = mp.plot(vtmp, ftmp, return_plot = True)
p1._Viewer__s["scale"] *= 1.2
reset(p1)
hand_full_rotation = {'s' : [3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : -2, 'α' : 1, 'β' : 0, 'γ' : 0, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5

p1.add_mesh(v_deformed, F)
add_text(p1, "CoR")
# p1.add_mesh(v_deformed + shft, F)
p1.add_mesh(U+2*shft, F)

add_text(p1, "DQS")


p1._Viewer__update_view()


p1.save("blending.html")

p1._Viewer__update_view()

play = iw.Play(
    value=0,
    min=0,
    max=40,
    step=1,
    interval=1000,
    description="Press play",
    disabled=False
)
slider = iw.IntSlider()
iw.jslink((play, 'value'), (slider, 'value'))
iw.HBox([play, slider])

hand_full_rotation = {'s' : [3], 'x': 0, 'y' : 0, 'z' : 0, 'w' : 1, 'α' : 0, 'β' : 0, 'γ' : 0}
p = mp.plot(V,F, return_plot=True)

def deform(x, hand_rot, plot):
    # reset(plot=plot)
    
    res = hand_rot | {'w': 1-x/10, 'α' : x/10}


    pos_f(**res, plt = plot)
    v_deformed = pos_f.deformer(OCoR)

    vT = np.array(pos_f_saver[:, :3])
    vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
    U = igl.dqs(V,W, vQ, vT)

    shft = np.zeros(V.shape)
    shft[:,1] = -.5

    # plot = mp.plot(V,F)
    # while len(plot._Viewer__objects.keys()) > 1:
    #     plot.remove_object(max(plot._Viewer__objects.keys()))
    
    
    
    # plot.add_mesh(v_deformed + shft, F)
    # plot.add_mesh(U+2*shft, F)
    # display(plot._renderer)
    plot._Viewer__update_view()
    

play = iw.Play(
    value=0,
    min=0,
    max=40,
    step=1,
    interval=1000,
    description="Press play",
    disabled=False
)
slider = iw.IntSlider()

p = mp.plot(V,F)
iw.interact(deform, x = slider, hand_rot=iw.fixed(hand_full_rotation), plot = iw.fixed(p))
iw.jslink((play, 'value'), (slider, 'value'))
iw.HBox([play, slider])

from pythreejs import *
from IPython.display import display
from math import pi

Geometry.from_geometry

# https://github.com/jupyter-widgets/pythreejs/blob/master/examples/Examples.ipynb

vertices = [
    [0, 0, 0],
    [0, 0, 1],
    [0, 1, 0],
    [0, 1, 1],
    [1, 0, 0],
    [1, 0, 1],
    [1, 1, 0],
    [1, 1, 1]
]

faces = [
    [0, 1, 3],
    [0, 3, 2],
    [0, 2, 4],
    [2, 6, 4],
    [0, 4, 1],
    [1, 4, 5],
    [2, 3, 6],
    [3, 7, 6],
    [1, 5, 3],
    [3, 5, 7],
    [4, 6, 5],
    [5, 6, 7]
]

vertexcolors = ['#000000', '#0000ff', '#00ff00', '#ff0000',
                '#00ffff', '#ff00ff', '#ffff00', '#ffffff']

# Map the vertex colors into the 'color' slot of the faces
faces = [f + [None, [vertexcolors[i] for i in f], None] for f in faces]

# Create the geometry:
cubeGeometry = Geometry(vertices=vertices,
    faces=faces,
    colors=vertexcolors)
# Calculate normals per face, for nice crisp edges:
cubeGeometry.exec_three_obj_method('computeFaceNormals')

# Create a mesh. Note that the material need to be told to use the vertex colors.
myobjectCube = Mesh(
    geometry=cubeGeometry,
    material=MeshLambertMaterial(vertexColors='VertexColors'),
    position=[-0.5, -0.5, -0.5],   # Center the cube
)

# Set up a scene and render it:
cCube = PerspectiveCamera(position=[3, 3, 3], fov=20,
                      children=[DirectionalLight(color='#ffffff', position=[-3, 5, 1], intensity=0.5)])
sceneCube = Scene(children=[myobjectCube, cCube, AmbientLight(color='#dddddd')])

rendererCube = Renderer(camera=cCube, background='black', background_opacity=1,
                        scene=sceneCube, controls=[OrbitControls(controlling=cCube)])

display(rendererCube)

# p1 = mp.plot(V,F)
# reset(p1)

hand_full_rotation = {'s' : [2,3], 'x': -.1, 'y' : 0.2, 'z' : 0.0, 'w' : .5, 'α' : -1, 'β' : -.5, 'γ' : 0, 'plt' : p1}
pos_f(**hand_full_rotation)
v_deformed = pos_f.deformer(OCoR)

vT = np.array(pos_f_saver[:, :3])
vQ = np.array(pos_f_saver[:, [4,5,6,3]], order='C')
U = igl.dqs(V,W, vQ, vT)

shft = np.zeros(V.shape)
shft[:,1] = -.5


# # p1.add_mesh(v_deformed + shft, F)
# p1.add_mesh(U+2*shft, F)

hV = V.astype('float32')
hF = F.astype('uint16').ravel()
handGeometry = BufferGeometry(attributes=dict(
    position=BufferAttribute(hV, normalized=False),
    index=BufferAttribute(hF, normalized=False)))

display(handGeometry)

V, F = igl.read_triangle_mesh('data/arm.obj')
C,BE,_,_,_,_ = igl.read_tgf('data/arm.tgf')
W = igl.read_dmat('data/arm-weights.dmat')
labels = np.array(range(BE.shape[0])).astype(int)

def apply_ocor(s,x,y,z, w, α, β, γ):
    clicked_edges = BE[np.isin(labels, s)]
    slices = []
    for e in clicked_edges:
        for v in e:
            if v not in slices:
                slices.append(v)
                
    r = Rotation.from_quat([α, β, γ, w])
    
    v_slice = C[slices] + np.array([[x,y,z]])
    
    C1 = C.copy()
    C1[slices] = r.apply(v_slice)

    q = np.repeat([[w,α, β, γ]], BE.shape[0], axis=0)
    t = np.repeat([[x,y,z]], BE.shape[0], axis=0)
#     t = np.zeros((BE.shape[0], 3))
#     for i in range(BE.shape[0]):
       
#         xi,yi,zi,wi,αi,βi,γi = pos_f_saver[i]
#         q[i,:] = [wi,αi,βi,γi]
#         t[i,:] = [xi, yi, zi]
    
    Q = quaternion.as_quat_array(q)
    
    OCoR = apply_weighted_rotation(V, W, Q, P, t)
    return OCoR

rotation = {'s' : [2,3], 'x': -.1, 'y' : 0.2, 'z' : 0.0, 'w' : .5, 'α' : -1, 'β' : -.5, 'γ' : 0}
nV = apply_ocor(**rotation)

hnV = nV.astype('float32')
handGeometry = BufferGeometry(attributes=dict(
    position=BufferAttribute(hV, normalized=False),
    index=BufferAttribute(hF, normalized=False)),
    morphAttributes=dict(
        position=[BufferAttribute(hnV, normalized=False)]))

display(handGeometry)

handGeometry.exec_three_obj_method('computeFaceNormals')

morphMesh = Mesh(handGeometry, MeshPhongMaterial(
    color='#ff3333', shininess=150, morphTargets=True))

pill_track = NumberKeyframeTrack(
    name='.morphTargetInfluences[0]', times=[0, 1.5, 3], values=[0, 2.5, 0])
pill_clip = AnimationClip(tracks=[pill_track])
pill_action = AnimationAction(AnimationMixer(morphMesh), pill_clip, morphMesh)

view_width = 600
view_height = 400

camera3 = PerspectiveCamera( position=[5, 3, 5], aspect=view_width/view_height)
scene3 = Scene(children=[morphMesh, camera3,
                         DirectionalLight(position=[3, 5, 1], intensity=0.6),
                         AmbientLight(intensity=0.5)])
renderer3 = Renderer(camera=camera3, scene=scene3,
                     controls=[OrbitControls(controlling=camera3)],
                     width=view_width, height=view_height)
display(renderer3, pill_action)

# push multiple iterations of the rotation and render through all of them and back

ball = Mesh(geometry=SphereGeometry(radius=1, widthSegments=32, heightSegments=24), 
            material=MeshLambertMaterial(color='red'),
            position=[2, 1, 0])

c = PerspectiveCamera(position=[0, 5, 5], up=[0, 1, 0],
                      children=[DirectionalLight(color='white', position=[3, 5, 1], intensity=0.5)])

scene = Scene(children=[ball, c, AmbientLight(color='#777777')])

renderer = Renderer(camera=c, 
                    scene=scene, 
                    controls=[OrbitControls(controlling=c)])
display(renderer)

# import time, math
# ball.material.color = '#4400dd'
# for i in range(1, 150, 2):
#     ball.scale = (i / 100.,) * 3
#     ball.position = [math.cos(i / 10.), math.sin(i / 50.), i / 100.]
#     time.sleep(.05)