import sys
from os.path import join
import time
import random
from pathlib import Path

import numpy as np

import matplotlib.pyplot as plt
import matplotlib.animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cbook
from matplotlib import cm
from matplotlib import animation



# ---------------------------------------------------------------------------
#
# Matplotlib drawing functions
#
# ---------------------------------------------------------------------------

# Generate random hex colors
def rhex():
  r = lambda: random.randint(0,255)
  return '#%02X%02X%02X' % (r(), r(), r())

 # line weight
def generate_3d_landmark_anim(lm, fp_out, num_frames=30, fps=12, dpi=72, size=(480,480),
  stroke_weight=2, mark_size=10, mark_type='.', bg_clr=(0,0,0), transparent=False):
  '''Generates animated 3D plot of face landmarks
  '''

  # convert opencv BGR numpy image to RGB
  bg_clr_hex = '#%02x%02x%02x' % bg_clr
  #mark_clr = '#%02x%02x%02x' % mark_clr
  
  # center x,y,z
  xmm = (np.min(lm[:,0]),np.max(lm[:,0]))
  ymm = (np.min(lm[:,1]),np.max(lm[:,1]))
  zmm = (np.min(lm[:,2]),np.max(lm[:,2]))
  
  # make copy of landmarks
  lm_orig = lm.copy()
  xmm = (np.min(lm_orig[:,0]),np.max(lm_orig[:,0]))
  ymm = (np.min(lm_orig[:,1]),np.max(lm_orig[:,1]))
  zmm = (np.min(lm_orig[:,2]),np.max(lm_orig[:,2]))
  
  # swap the y and z components to improve 3d rotation angles for matplotlib
  lm = np.zeros_like(lm_orig).astype(np.uint8)
  for i,p in enumerate(lm_orig):
    x,y,z = p
    lm[i] = np.array([x - xmm[0], z - zmm[0], y - ymm[0]])
  
  # Create plot
  figsize = (size[0]/dpi, size[1]/dpi )
  fig = plt.figure(figsize=figsize, dpi=dpi)  # frameon=False
  fig.tight_layout()
  # remove whitespace in matplotlib
  fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
  ax = fig.add_subplot(111, projection='3d')
  ax.set_facecolor(bg_clr_hex) # background color
  
  xscale, yscale, zscale = (1.2, 1.0, 1.0)
  
  # scatter plot the dots

  # jaw line
  mark_clr = '#%02x%02x%02x' % (0,255,0)  # green
  ax.plot3D(lm[:17,0]*1.2,lm[:17,1], lm[:17,2],
            marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)

  # stage-right eyebrow
  mark_clr = '#%02x%02x%02x' % (255,0,0)  # green
  ax.plot3D(lm[17:22,0]*1.2,lm[17:22,1],lm[17:22,2],
            marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # stage-left eyebrow
  mark_clr = '#%02x%02x%02x' % (255,255,0)  # yellow
  ax.plot3D(lm[22:27,0]*1.2,lm[22:27,1],lm[22:27,2], 
            marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # nose ridge
  mark_clr = '#%02x%02x%02x' % (0,0,255)  # blue
  ax.plot3D(lm[27:31,0]*1.2,lm[27:31,1],lm[27:31,2],
            marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # nose-bottom
  mark_clr = '#%02x%02x%02x' % (255,0,255)  # magenta
  ax.plot3D(lm[31:36,0]*1.2,lm[31:36,1],lm[31:36,2],
            marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # stage-left eye
  mark_clr = '#%02x%02x%02x' % (0,255,255)  # cyan
  px, py, pz = lm[36:42,0]*1.2,lm[36:42,1],lm[36:42,2]
  px = np.append(px, lm[36,0]*1.2)
  py = np.append(py, lm[36,1])
  pz = np.append(pz, lm[36,2])
  ax.plot3D(px, py, pz, marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # stage-right eye
  mark_clr = '#%02x%02x%02x' % (255,255,255)  # white
  px, py, pz = lm[42:48,0]*1.2,lm[42:48,1],lm[42:48,2]
  px = np.append(px, lm[42,0]*1.2)
  py = np.append(py, lm[42,1])
  pz = np.append(pz, lm[42,2])
  ax.plot3D(px, py, pz, marker=mark_type, markersize=mark_size, color=mark_clr,linewidth=stroke_weight)
  
  # mouth
  mark_clr = '#%02x%02x%02x' % (255,125,0)  # orange?
  px, py, pz = lm[48:,0]*1.2,lm[48:,1],lm[48:,2]
  px = np.append(px, lm[48,0]*1.2)
  py = np.append(py, lm[48,1])
  pz = np.append(pz, lm[48,2])
  ax.plot3D(px, py, pz, marker=mark_type, markersize=mark_size, color=mark_clr, linewidth=stroke_weight)
  
  #rh = '#00ff00'  # edge color
  #ax.scatter(lm[:,0]*xscale,lm[:,1]*yscale,lm[:,2]*zscale, c=rh, alpha=1.0, s=35, edgecolor=rh)
  #ax.scatter(lm[:,0]*xscale,lm[:,1]*yscale,lm[:,2]*zscale, c=rh, alpha=1.0, s=1)
  
  # center center x,y,z points
  cx = ((xmm[0] - xmm[1]) // 2) + xmm[1]
  cy = ((ymm[1] - ymm[0]) // 2) + ymm[0]
  cz = ((zmm[1] - zmm[0]) // 2) + zmm[0]
  
  # remove ticks
  ax.set_xticks([])
  ax.set_yticks([])
  ax.set_zticks([])
    
  # remove axis
  ax.set_frame_on(False)
  ax.set_axis_off()

  # set initial plot view
  ax.view_init(elev=120., azim=70.)

  # rotation increments: from 0 to 360 in num_frames
  phi = np.linspace(0, 2*np.pi, num_frames)

  # animation instruction
  def update(phi):
    ax.view_init(180,phi*180./np.pi)
  
  ani = matplotlib.animation.FuncAnimation(fig, update, frames=phi)
  
  savefig_kwargs = {'pad_inches': 0, 'transparent': transparent}
  ani.save(fp_out, writer='imagemagick', fps=fps, savefig_kwargs=savefig_kwargs)