import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision
from mediapipe.framework.formats import landmark_pb2
from mediapipe import solutions
import numpy as np
import time
import cv2
import argparse
import os
import math

# modified in gradio

from mp_constants import *
from mp_utils import divide_line_to_points,points_to_bbox,expand_bbox

import logging

# for share lib,TODO make module
#import sys
#sys.path.append("C:\\Users\\owner\\Documents\\pythons\\glibvision")
from glibvision.glandmark_utils import bbox_to_glandmarks,convert_to_landmark_group_json
from glibvision.cv2_utils import draw_bbox,plot_points,set_plot_text

def parse_arguments():
  """
  引数
  
  """
  parser = argparse.ArgumentParser(
          description="draw 68 points"
      )
  parser.add_argument(
          "--input_file","-i",required=True,help="Input file"
      )
  parser.add_argument(
          "--model_path","-m",default="face_landmarker.task",help="model path"
      )
  parser.add_argument(
          "--save_glandmark","-g",action="store_true",help="save godot-landmark json"
      )
  parser.add_argument(
          "--save_group_landmark","-landmark",action="store_true",help="save group-landmark json"
      )
  return parser.parse_args()





def draw_landmarks_on_image(rgb_image, detection_result,draw_number,font_scale,text_color,dot_size,dot_color,line_size,line_color,box_size,box_color):
  #print(f"dot_size={dot_size},dot_color={dot_color},line_size={line_size},line_color={line_color}")
  image_width,iamge_height = rgb_image.size
  face_landmarks_list = detection_result.face_landmarks
  annotated_image = np.copy(rgb_image)

  def get_cordinate(index):
    x=face_landmarks_list[0][index].x
    y=face_landmarks_list[0][index].y
    return x,y
    
  def get_distance(x1,y1,x2,y2):
    return math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
  
  def get_centers():
    center_indices =[
    #(POINT_LEFT_HEAD_OUTER,POINT_RIGHT_HEAD_OUTER,POINT_FOREHEAD_TOP),
    #(POINT_LEFT_HEAD_OUTER,POINT_RIGHT_HEAD_OUTER,POINT_CHIN_BOTTOM),
    [POINT_NOSE_CENTER_MIDDLE],
    #[POINT_LOWER_LIP_CENTER_BOTTOM]
    #(POINT_UPPER_LIP_CENTER_BOTTOM,POINT_LOWER_LIP_CENTER_TOP)
    ]
    centers = []
    for indices in center_indices:
      total_x = 0
      total_y = 0
      for index in indices:
        x,y = get_cordinate(index)
        total_x+=x
        total_y+=y
      centers.append ((total_x/len(indices),total_y/len(indices)))
    return centers
  
  centers = get_centers()
  for center in centers:
    center_x,center_y = center

    pt = int(center_x*image_width),int(center_y*iamge_height)
    
    #cv2.circle(annotated_image,pt,20,(0,0,255),-1)

  def get_closed_center(x,y):
    closed = None
    closed_distance = 0
    for center in centers:
      distance = get_distance(center[0],center[1],x,y)
      if closed == None:
        closed = center
        closed_distance = distance
      else:
        if distance<closed_distance:
          closed_distance = distance
          closed = center
    return closed
  

  #landmark is [index-upper,index-lower]
  def get_mean_point(landmark,width=image_width,height=iamge_height):
    xs=[]
    ys=[]
    for index in landmark:
      x,y = get_cordinate(index) #inner cordinate
      xs.append(x)
      ys.append(y)

    return int(np.mean(xs)*width),int(np.mean(ys)*height)
    
  def get_cordinate_point(landmark,width=image_width,height=iamge_height):
      point = get_cordinate(landmark)

      return int(point[0]*width),int(point[1]*height)
  # TODO rename and explain this is for contour choose most outer point
  def get_point(landmark,width=image_width,height=iamge_height):
    xs=[]
    ys=[]
    

    def get_outer_point(indexes):
      outer_point = None
      max_distance = None
      if len(indexes) == 0:
        return None
      
      ratio = 0.5
      x,y = get_cordinate(indexes[-1]) #on contour 3 lines outer,center,inner cordinate

      #x,y = get_cordinate(indexes[0])
      center_x,center_y = get_closed_center(x,y)
      x-=(center_x-x)*ratio
      y-=(center_y-y)*ratio

      outer_x = x
      outer_y = y

      for index in indexes:
        x,y = get_cordinate(index)
        
        distance = get_distance(outer_x,outer_y,x,y)
        #print(f"{distance} index={index} x={x},y={y}")
        if outer_point == None:
          outer_point = (x,y)
          max_distance = distance
        else:
          if distance<max_distance:
            outer_point = (x,y)
      return outer_point



    for group in landmark:
      outer_point = get_outer_point(group)
      xs.append(outer_point[0])
      ys.append(outer_point[1])
      

    return int(np.mean(xs)*width),int(np.mean(ys)*height)

  # Loop through the detected faces to visualize.
  for idx in range(len(face_landmarks_list)):
    face_landmarks = face_landmarks_list[idx]

    # Draw the face landmarks. #something change format
    face_landmarks_proto = landmark_pb2.NormalizedLandmarkList()
    face_landmarks_proto.landmark.extend([
      landmark_pb2.NormalizedLandmark(x=landmark.x, y=landmark.y, z=landmark.z) for landmark in face_landmarks
    ])


    def draw_sets(draw_set,color=(0,255,0)):
        solutions.drawing_utils.draw_landmarks(
        image=annotated_image,
        landmark_list=face_landmarks_proto,
        connections=draw_set,
        landmark_drawing_spec=None,
        connection_drawing_spec=mp.solutions.drawing_styles.DrawingSpec(color=color, thickness=1 ))
    def draw_triangle(index1,index2,index3):
       draw_sets({(index1,index2),(index2,index3),(index3,index1)})

    def draw_lines(array,color=(0,0,128)):
      my_set = set()
      for i in range(len(array)-1):
         v = (array[i],array[i+1])
         my_set.add(v)
      draw_sets(my_set,color)

    def convert_to_box(face_landmarks_list,indices,w=1024,h=1024):
      x1=0
      y1=0
      x2=w
      y2=h
      for index in indices:
        x=min(w,max(0,(face_landmarks_list[0][index].x*w)))
        y=min(h,max(0,(face_landmarks_list[0][index].y*h)))
        if x>x1:
          x1=x
        if y>y1:
          y1=y
        
        if x<x2:
          x2=x
        if y<y2:
          y2=y

      return [x1,y1,x2-x1,y2-y1]
          
       
    my_set ={(362,382),(382,398),(398,362)}
    my_set = mp.solutions.face_mesh.FACEMESH_RIGHT_EYE

    #mediapipe to 5point
    """
    draw_triangle(362,382,398)
    draw_triangle(173,133,155)
    draw_triangle(33,246,7)
    draw_triangle(249,263,466)

    draw_triangle(94,2,164)

    draw_triangle(61,76,61)
    draw_triangle(291,306,291)
   
    draw_lines([17,18,200,199,175,152])

    draw_lines([127,234,93,132,58,172,136,150,149,176,148,152],(255,0,0))
    #draw_lines([127,234,132,172,150,176,152],(0,0,255))
    """

    #
    #draw_lines([9,107])
    """
    draw_lines([148,171,208])
    draw_lines([176,140])
    draw_lines([149,170,211])
    draw_lines([150,169,])
   
    draw_lines([150,169])
    draw_lines([136,135,214])
    draw_lines([172,138,192])
    draw_lines([58,215])
    draw_lines([132,177,147])
    draw_lines([58,215,213])
    draw_lines([93,137,123])
    #draw_lines([234,227])
    #draw_lines([127,34,143])

    """
    #draw_lines([378,288,356,251,151,21,127,58,150,152])
    #draw_lines(LINE_RIGHT_CONTOUR_OUTER_EYE_TO_CHIN)
    #draw_lines(LINE_RIGHT_CONTOUR_EYE_TO_CHIN)
    #draw_lines(LINE_RIGHT_CONTOUR_INNER_EYE_TO_CHIN,(0,255,0))
    """
    draw_lines(LINE_RIGHT_CONTOUR_0)
    draw_lines(LINE_RIGHT_CONTOUR_1)
    draw_lines(LINE_RIGHT_CONTOUR_2)
    draw_lines(LINE_RIGHT_CONTOUR_3)
    draw_lines(LINE_RIGHT_CONTOUR_4)
    draw_lines(LINE_RIGHT_CONTOUR_5)
    draw_lines(LINE_RIGHT_CONTOUR_6)
    draw_lines(LINE_RIGHT_CONTOUR_7)
    draw_lines(LINE_RIGHT_CONTOUR_8)
    draw_lines(LINE_RIGHT_CONTOUR_9)
    draw_lines(LINE_RIGHT_CONTOUR_10)
    draw_lines(LINE_RIGHT_CONTOUR_11)

    
    draw_lines(LINE_LEFT_CONTOUR_1)
    draw_lines(LINE_LEFT_CONTOUR_2)
    draw_lines(LINE_LEFT_CONTOUR_3)
    draw_lines(LINE_LEFT_CONTOUR_4)
    draw_lines(LINE_LEFT_CONTOUR_5)
    draw_lines(LINE_LEFT_CONTOUR_6)
    draw_lines(LINE_LEFT_CONTOUR_7)
    draw_lines(LINE_LEFT_CONTOUR_8)
    draw_lines(LINE_LEFT_CONTOUR_9)
    draw_lines(LINE_LEFT_CONTOUR_10)
    draw_lines(LINE_LEFT_CONTOUR_11)
    #draw_lines(LINE_LEFT_CONTOUR_12)
    """

    #draw_lines(LINE_RIGHT_CONTOUR_6,(255,0,0))
    
    def get_eye_brow_points(landmarks):
      result_points= []
      for landmark in landmarks:
        point=get_mean_point(landmark)
        result_points.append(point)

      return result_points
    
    def get_mean_points(landmarks):
      result_points= []
      for landmark in landmarks:
        point=get_mean_point(landmark)
        result_points.append(point)

      return result_points
    
    def get_divided_points(landmarks,divided=3):
      result_points= []
      landmark_points = []
      for landmark in landmarks:
        if isinstance(landmark, int):
          pt=get_cordinate_point(landmark)
        else:
          pt =get_mean_point(landmark)
        landmark_points.append(pt)

      divided_points = divide_line_to_points(landmark_points,divided)

      #print(centers[0][0]*1024,",",centers[0][1]*1024)
      return divided_points
    
    
    def get_half_contour(landmarks):
      result_points= []
      landmark_points = []
      for landmark in landmarks:
        pt =get_point(landmark)
        landmark_points.append(pt)

      divided_points = divide_line_to_points(landmark_points,8)#9
      #for pt in divided_points:
      #  cv2.circle(annotated_image,pt,3,(255,0,0),-1)
      #  result_points.append((pt[0],pt[1]))

      #print(centers[0][0]*1024,",",centers[0][1]*1024)
      return divided_points
    
    right_landmarks =[
      #[LANDMARK_68_CONTOUR_5]
      [LANDMARK_68_CONTOUR_1],[LANDMARK_68_CONTOUR_2_PART1,LANDMARK_68_CONTOUR_2_PART2],[LANDMARK_68_CONTOUR_3],[LANDMARK_68_CONTOUR_4],[LANDMARK_68_CONTOUR_5],[LANDMARK_68_CONTOUR_6_PART1,LANDMARK_68_CONTOUR_6_PART2],[LANDMARK_68_CONTOUR_7],[LANDMARK_68_CONTOUR_8_PART1,LANDMARK_68_CONTOUR_8_PART2],[LANDMARK_68_CONTOUR_9],
     
    ]
    contour_right_points=get_half_contour(right_landmarks)

    left_landmarks =[
     [LANDMARK_68_CONTOUR_9], [LINE_LEFT_CONTOUR_1], [LINE_LEFT_CONTOUR_2], [LINE_LEFT_CONTOUR_3], [LINE_LEFT_CONTOUR_4],[LINE_LEFT_CONTOUR_5],[LINE_LEFT_CONTOUR_6],[LINE_LEFT_CONTOUR_7],[LINE_LEFT_CONTOUR_8],[LINE_LEFT_CONTOUR_9],[LINE_LEFT_CONTOUR_10],[LINE_LEFT_CONTOUR_11]
    ]
    contour_left_points=get_half_contour(left_landmarks)

    set_plot_text(draw_number,font_scale,text_color) # for reset
    plot_points(annotated_image,contour_right_points+contour_left_points[1:],False,dot_size,dot_color,line_size,line_color)

    right_eye_brow_points=get_eye_brow_points([
      LANDMARK_68_RIGHT_EYEBROW_18,LANDMARK_68_RIGHT_EYEBROW_19,LANDMARK_68_RIGHT_EYEBROW_20,LANDMARK_68_RIGHT_EYEBROW_21,LANDMARK_68_RIGHT_EYEBROW_22
    ])
    plot_points(annotated_image,right_eye_brow_points,False,dot_size,dot_color,line_size,line_color)
    left_eye_brow_points=get_eye_brow_points([
      LANDMARK_68_LEFT_EYEBROW_23,LANDMARK_68_LEFT_EYEBROW_24,LANDMARK_68_LEFT_EYEBROW_25,LANDMARK_68_LEFT_EYEBROW_26,LANDMARK_68_LEFT_EYEBROW_27
    ])
    plot_points(annotated_image,left_eye_brow_points,False,dot_size,dot_color,line_size,line_color)

    vertical_nose_points = get_divided_points([LANDMARK_68_VERTICAL_NOSE_28,LANDMARK_68_VERTICAL_NOSE_29,LANDMARK_68_VERTICAL_NOSE_30,LANDMARK_68_VERTICAL_NOSE_31],3)
    plot_points(annotated_image,vertical_nose_points,False,dot_size,dot_color,line_size,line_color)

    horizontal_nose_points = get_mean_points([LANDMARK_68_HORIZONTAL_NOSE_32,LANDMARK_68_HORIZONTAL_NOSE_33,LANDMARK_68_HORIZONTAL_NOSE_34,LANDMARK_68_HORIZONTAL_NOSE_35,LANDMARK_68_HORIZONTAL_NOSE_36])
    plot_points(annotated_image,horizontal_nose_points,False,dot_size,dot_color,line_size,line_color)

    right_upper_eye_points = get_divided_points(LINE_RIGHT_UPPER_MIXED_EYE2,3)
    right_lower_eye_points = get_divided_points(LINE_RIGHT_LOWER_MIXED_EYE,3)
    #right_eye_points = right_upper_eye_points+right_lower_eye_points # first and last is same as above
    right_eye_points = right_upper_eye_points+right_lower_eye_points[1:-1]
    plot_points(annotated_image,right_eye_points,True,dot_size,dot_color,line_size,line_color)
    
    #draw_lines(LINE_RIGHT_LOWER_OUTER_EYE,(0,255,0))
    #draw_lines(LINE_RIGHT_LOWER_INNER_EYE,(0,255,0))
    #draw_lines(LINE_RIGHT_UPPER_OUTER_EYE,(0,255,0))
    #draw_lines(LINE_RIGHT_UPPER_INNER_EYE,(0,255,0))

    left_upper_eye_points = get_divided_points(LINE_LEFT_UPPER_MIXED_EYE2,3)
    left_lower_eye_points = get_divided_points(LINE_LEFT_LOWER_MIXED_EYE,3)
    #left_eye_points = left_upper_eye_points+left_lower_eye_points# first and last is same as above
    left_eye_points = left_upper_eye_points+left_lower_eye_points[1:-1]
    plot_points(annotated_image,left_eye_points,True,dot_size,dot_color,line_size,line_color)
    # first and last is same as above

    #draw_lines(LINE_LEFT_LOWER_OUTER_EYE,(0,255,0))
    #draw_lines(LINE_LEFT_LOWER_INNER_EYE,(0,255,0))
    #draw_lines(LINE_LEFT_UPPER_OUTER_EYE,(0,255,0))
    #draw_lines(LINE_LEFT_UPPER_INNER_EYE,(0,255,0))

    left_upper_outer_lip_points = get_divided_points(LINE_RIGHT_UPPER_OUTER_LIP,3)
    right_upper_outer_lip_points = get_divided_points(LINE_LEFT_UPPER_OUTER_LIP,3)
    upper_outer_lip_points = left_upper_outer_lip_points+right_upper_outer_lip_points[1:]# first and last is same as above
    #plot_points(annotated_image,upper_outer_lip_points)
    upper_outer_lip_points = get_mean_points([LANDMARK_68_UPPER_OUTER_LIP_49,LANDMARK_68_UPPER_OUTER_LIP_50,LANDMARK_68_UPPER_OUTER_LIP_51,LANDMARK_68_UPPER_OUTER_LIP_52,LANDMARK_68_UPPER_OUTER_LIP_53,LANDMARK_68_UPPER_OUTER_LIP_54,LANDMARK_68_UPPER_OUTER_LIP_55])
    #plot_points(annotated_image,upper_outer_lip_points)

    lower_outer_lip_points = get_mean_points([LANDMARK_68_UPPER_OUTER_LIP_55,LANDMARK_68_LOWER_OUTER_LIP_56,LANDMARK_68_LOWER_OUTER_LIP_57,LANDMARK_68_LOWER_OUTER_LIP_58,LANDMARK_68_LOWER_OUTER_LIP_59,LANDMARK_68_LOWER_OUTER_LIP_60,LANDMARK_68_UPPER_OUTER_LIP_49])
    outer_lip_points = upper_outer_lip_points+lower_outer_lip_points[1:-1]
    plot_points(annotated_image,outer_lip_points,True,dot_size,dot_color,line_size,line_color)

    upper_inner_lip_points = get_mean_points([LANDMARK_68_UPPER_INNER_LIP_61,LANDMARK_68_UPPER_INNER_LIP_62,LANDMARK_68_UPPER_INNER_LIP_63,LANDMARK_68_UPPER_INNER_LIP_64,LANDMARK_68_UPPER_INNER_LIP_65])
    #plot_points(annotated_image,upper_inner_lip_points)

    lower_inner_lip_points = get_mean_points([LANDMARK_68_UPPER_INNER_LIP_65,LANDMARK_68_LOWER_INNER_LIP_66,LANDMARK_68_LOWER_INNER_LIP_67,LANDMARK_68_LOWER_INNER_LIP_68,LANDMARK_68_UPPER_INNER_LIP_61])
    inner_lip_points = upper_inner_lip_points+lower_inner_lip_points[1:-1]
    plot_points(annotated_image,inner_lip_points,True,dot_size,dot_color,line_size,line_color)

    landmark_points = contour_right_points+contour_left_points[1:]
    landmark_points += right_eye_brow_points + left_eye_brow_points
    landmark_points += vertical_nose_points + horizontal_nose_points
    landmark_points += right_eye_points + left_eye_points
    landmark_points += outer_lip_points + inner_lip_points

    #plot_points(annotated_image,landmark_points,20) # for debug
    bbox = points_to_bbox(landmark_points)
    bbox = expand_bbox(bbox,5,7,5,5)

    draw_bbox(annotated_image,bbox,box_color,box_size)
    #draw_lines([POINT_LEFT_HEAD_OUTER_EX,POINT_LEFT_EYE_OUTER_EX,POINT_LEFT_MOUTH_OUTER_EX,POINT_LEFT_CHIN_OUTER,POINT_CHIN_BOTTOM])
    #draw_lines([LANDMARK_68_CONTOUR_1,LANDMARK_68_CONTOUR_2,LANDMARK_68_CONTOUR_3,LANDMARK_68_CONTOUR_4,LANDMARK_68_CONTOUR_5,LANDMARK_68_CONTOUR_6,LANDMARK_68_CONTOUR_7,LANDMARK_68_CONTOUR_8,LANDMARK_68_CONTOUR_9])
    """solutions.drawing_utils.draw_landmarks(
        image=annotated_image,
        landmark_list=face_landmarks_proto,
        connections=mp.solutions.face_mesh.FACEMESH_LEFT_EYE,#FACE_OVAL
        landmark_drawing_spec=None,
        connection_drawing_spec=mp.solutions.drawing_styles
        .get_default_face_mesh_contours_style())"""

    """solutions.drawing_utils.draw_landmarks(
        image=annotated_image,
        landmark_list=face_landmarks_proto,
        connections=mp.solutions.face_mesh.FACEMESH_CONTOURS,# mix all
        landmark_drawing_spec=None,
        connection_drawing_spec=mp.solutions.drawing_styles
        .get_default_face_mesh_contours_style())  
     """
    """solutions.drawing_utils.draw_landmarks(
        image=annotated_image,
        landmark_list=face_landmarks_proto,
        connections=mp.solutions.face_mesh.FACEMESH_IRISES,
          landmark_drawing_spec=None,
          connection_drawing_spec=mp.solutions.drawing_styles
          .get_default_face_mesh_iris_connections_style()) """

  return annotated_image,bbox,landmark_points


if __name__ == "__main__":
  args = parse_arguments()
  input_file = args.input_file

#file checks
  if not os.path.isfile(input_file):
    print(f"input is not file '{input_file}'")
    exit(0)

  model_path = args.model_path

  BaseOptions = mp.tasks.BaseOptions
  FaceLandmarker = mp.tasks.vision.FaceLandmarker
  FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
  VisionRunningMode = mp.tasks.vision.RunningMode

  options = FaceLandmarkerOptions(
      base_options=BaseOptions(model_asset_path=model_path),
      running_mode=VisionRunningMode.IMAGE
      ,min_face_detection_confidence=0, min_face_presence_confidence=0
      )


  with FaceLandmarker.create_from_options(options) as landmarker:
      
      start = time.time()
      mp_image = mp.Image.create_from_file(input_file)
      face_landmarker_result = landmarker.detect(mp_image)
      detect_time = time.time()-start
      print(detect_time)
      
      annotated_image,bbox,landmark_points = draw_landmarks_on_image(mp_image.numpy_view(), face_landmarker_result)
      #print(annotated_image)
      #annotated_image=cv2.resize(annotated_image, (800, 800))
      annotated_image=cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
      cv2.imwrite(input_file.replace(".jpg","_la68.jpg"),annotated_image)

      if args.save_glandmark:
        parent_path,file = os.path.split(input_file)
        glandmark = bbox_to_glandmarks(file,bbox,landmark_points)
        glandmark_path = input_file.replace(".jpg",f".json")
        if os.path.exists(glandmark_path):
            print(f"glandmark exist skipped {glandmark_path}")
        else:
            import json
            with open(glandmark_path,"w") as f:
                json.dump(glandmark,f)

      # _landmark.json always overwrite because not design for edit
      if args.save_group_landmark:
        result=convert_to_landmark_group_json(landmark_points)
        total = 0
        for key in result[0].keys():
          total += len(result[0][key])

        print(total)
        import json
        group_landmark_path = input_file.replace(".jpg",f"_landmark.json")
        with open(group_landmark_path,"w") as f:
              json.dump(result,f)

    #cv2.imshow("image",)
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()