import ast
from pathlib import Path, PosixPath

import cv2
import dlib
import imutils
import numpy as np
from imutils.face_utils.facealigner import FaceAligner
from retinaface import RetinaFace

# ROOT_DIR = Path(__file__).parent


# Detect face in image and return bounding box of face
#
def detect_face_bb(img_path: str) -> tuple | None:
    faces = RetinaFace.detect_faces(img_path=img_path)

    for key in faces:
        if not key:
            return None
        bb = faces[key]
        (x1, y1, x2, y2) = (
            bb["facial_area"][0],
            bb["facial_area"][1],
            bb["facial_area"][2],
            bb["facial_area"][3],
        )
        return (x1, y1, x2, y2)


# Get path to pretrained landmarks detector
#
def get_predictor(path: str):
    landmark_detector = dlib.shape_predictor(
        (Path(path) / "shape_predictor_68_face_landmarks.dat").as_posix()
    )
    return landmark_detector


# Get lanmarks from image and return them as array
#
def get_landmarks(img: np.ndarray) -> np.ndarray | None:
    detector = dlib.get_frontal_face_detector()
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    rects = detector(gray, 1)
    if len(rects) == 0:
        return None
    landmark_detector = get_predictor("../data")
    dlib_landmarks = landmark_detector(gray, rects[0])
    landmarks_to_array = transform_landmarks_to_array(dlib_landmarks)

    return landmarks_to_array


# Gets a DLIB landmarks object and returns a [68,2] numpy array with the
# landmark coordinates.
def transform_landmarks_to_array(landmarks_dlib_obj) -> np.ndarray:
    # Gets a DLIB landmarks object and returns a [68,2] numpy array with the
    # landmark coordinates.

    landmark_array = np.zeros([68, 2])

    for i in range(68):
        landmark_array[i, 0] = landmarks_dlib_obj.part(i).x
        landmark_array[i, 1] = landmarks_dlib_obj.part(i).y

    return landmark_array  # numpy array


# Draw detected landmarks into image.
#
def draw_landmarks(img: np.ndarray, landmarks_array: np.ndarray) -> np.ndarray:
    for n in range(0, 68):
        draw_landmark(img, landmarks_array, n)

    return img


# Draw landmark into image based on its coordinates.
#
def draw_landmark(img: np.ndarray, landmarks_array: np.ndarray, n: int) -> np.ndarray:
    x = int(landmarks_array[n, 0])
    y = int(landmarks_array[n, 1])

    cv2.circle(img, (x, y), 2, (21, 252, 252), -1)

    return img


# Draw defined landmark into image.
#
def defined_landmarks(img: np.ndarray, landmarks: np.ndarray) -> np.ndarray:
    draw_landmark(img, landmarks, 48)
    draw_landmark(img, landmarks, 54)

    draw_landmark(img, landmarks, 39)
    draw_landmark(img, landmarks, 42)

    draw_landmark(img, landmarks, 36)
    draw_landmark(img, landmarks, 45)

    draw_landmark(img, landmarks, 0)
    draw_landmark(img, landmarks, 16)

    draw_landmark(img, landmarks, 31)
    draw_landmark(img, landmarks, 35)

    draw_landmark(img, landmarks, 27)
    draw_landmark(img, landmarks, 33)

    draw_landmark(img, landmarks, 27)
    draw_landmark(img, landmarks, 8)

    draw_landmark(img, landmarks, 33)
    draw_landmark(img, landmarks, 8)

    draw_landmark(img, landmarks, 62)
    draw_landmark(img, landmarks, 8)

    draw_landmark(img, landmarks, 4)
    draw_landmark(img, landmarks, 12)

    return img


# Scale image by converting into grayscale, rotation and zoom
#
def scale_image(img: np.ndarray, zoom: float) -> np.ndarray | None:
    # resize image
    image = imutils.resize(img, width=800)
    # convert input image into grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # detect face in image
    bb = detect_face_bb(image)

    if bb is None:
        return None

    # convert bounding box to dlib rectangle
    rect = dlib.rectangle(left=bb[0], top=bb[1], right=bb[2], bottom=bb[3])

    # use face aligner to rotate and scale image
    predictor = get_predictor("../data")
    fa = FaceAligner(predictor, desiredLeftEye=(zoom, zoom), desiredFaceWidth=256)
    faceAligned = fa.align(image, gray, rect)

    # return scaled image
    return faceAligned


# Parse id and age from image name and return them as tuple
#
def parse_image_name(img_name: PosixPath) -> tuple:
    id, *age = img_name.with_suffix("").name.split("_")
    id = int(id)
    age = float(".".join(age))
    return id, age


# Convert values in string into numpy array
#
def read_np_array(array_string) -> np.ndarray:
    array_string = ",".join(array_string.replace("[  ", "[").replace("[ ", "[").split())
    return np.array(ast.literal_eval(array_string))


# Normalization function
#
def norm_func(x, a, b, c, d) -> float:
    return a / (1 + np.exp(-x * c + d)) + b


# Calculate all measurements as euclidean distance of 2 landmarks
#
def get_measurements(landmarks: np.ndarray) -> dict:
    measurements = {}
    measurements["en_en"] = np.linalg.norm(landmarks[39] - landmarks[42])
    measurements["ex_ex"] = np.linalg.norm(landmarks[36] - landmarks[45])
    measurements["al_al"] = np.linalg.norm(landmarks[31] - landmarks[35])
    measurements["ch_ch"] = np.linalg.norm(landmarks[48] - landmarks[54])
    measurements["go_go"] = np.linalg.norm(landmarks[4] - landmarks[12])
    measurements["zy_zy"] = np.linalg.norm(landmarks[0] - landmarks[16])
    measurements["n_gna"] = np.linalg.norm(landmarks[27] - landmarks[8])
    measurements["n_sn"] = np.linalg.norm(landmarks[27] - landmarks[33])
    measurements["sn_gna"] = np.linalg.norm(landmarks[33] - landmarks[8])
    measurements["sto_gna"] = (
        np.linalg.norm(landmarks[62] - landmarks[8])
        - np.linalg.norm(landmarks[62] - landmarks[66]) / 2
    )

    return measurements


# Calculate landmarks coordinates at predicted ages.
#
def age_progress(count, plot_age, input_points, shift, predict_measurement):
    final_df = []
    for i in range(count):
        points_df = []
        points_df.append(input_points)
        for age in plot_age:
            future_points = input_points.copy()
            future_points.loc["zy_l"]["x"] = input_points.loc["zy_l"]["x"] - (
                shift[i].loc[age]["zy-zy"] / 2
            )
            future_points.loc["zy_r"]["x"] = input_points.loc["zy_r"]["x"] + (
                shift[i].loc[age]["zy-zy"] / 2
            )

            future_points.loc["en_l"]["x"] = input_points.loc["en_l"]["x"] - (
                shift[i].loc[age]["en-en"] / 2
            )
            future_points.loc["en_r"]["x"] = input_points.loc["en_r"]["x"] + (
                shift[i].loc[age]["en-en"] / 2
            )

            future_points.loc["ex_l"]["x"] = input_points.loc["ex_l"]["x"] - (
                shift[i].loc[age]["ex-ex"] / 2
            )
            future_points.loc["ex_r"]["x"] = input_points.loc["ex_r"]["x"] + (
                shift[i].loc[age]["ex-ex"] / 2
            )

            future_points.loc["al_l"]["x"] = input_points.loc["al_l"]["x"] - (
                shift[i].loc[age]["al-al"] / 2
            )
            future_points.loc["al_r"]["x"] = input_points.loc["al_r"]["x"] + (
                shift[i].loc[age]["al-al"] / 2
            )

            future_points.loc["ch_l"]["x"] = input_points.loc["ch_l"]["x"] - (
                shift[i].loc[age]["ch-ch"] / 2
            )
            future_points.loc["ch_r"]["x"] = input_points.loc["ch_r"]["x"] + (
                shift[i].loc[age]["ch-ch"] / 2
            )

            future_points.loc["go_l"]["x"] = input_points.loc["go_l"]["x"] - (
                shift[i].loc[age]["go-go"] / 2
            )
            future_points.loc["go_r"]["x"] = input_points.loc["go_r"]["x"] + (
                shift[i].loc[age]["go-go"] / 2
            )

            # vertical growth of nose
            future_points.loc["sn"]["y"] = (
                input_points.loc["sn"]["y"] + shift[i].loc[age]["n-sn"]
            )
            future_points.loc["al_r"]["y"] = (
                input_points.loc["al_r"]["y"] + shift[i].loc[age]["n-sn"]
            )
            future_points.loc["al_l"]["y"] = (
                input_points.loc["al_l"]["y"] + shift[i].loc[age]["n-sn"]
            )

            future_points.loc["go_r"]["y"] = (
                input_points.loc["go_r"]["y"] + shift[i].loc[age]["n-sn"]
            )
            future_points.loc["go_l"]["y"] = (
                input_points.loc["go_l"]["y"] + shift[i].loc[age]["n-sn"]
            )

            future_points.loc["gna"]["y"] = (
                input_points.loc["gna"]["y"] + shift[i].loc[age]["n-gna"]
            )

            future_points.loc["sto"]["y"] = (
                future_points.loc["gna"]["y"]
                - predict_measurement[i].loc[age]["sto-gna"]
            )

            future_points.loc["ch_r"]["y"] = input_points.loc["ch_r"]["y"] + (
                future_points.loc["sto"]["y"] - input_points.loc["sto"]["y"]
            )
            future_points.loc["ch_l"]["y"] = input_points.loc["ch_l"]["y"] + (
                future_points.loc["sto"]["y"] - input_points.loc["sto"]["y"]
            )

            points_df.append(future_points)

        final_df.append(points_df)
    return final_df
