from functools import lru_cache
from typing import List, Sequence, Tuple

import cv2
import numpy
from cv2.typing import Size

from facefusion.types import Anchors, Angle, BoundingBox, Distance, FaceDetectorModel, FaceLandmark5, FaceLandmark68, Mask, Matrix, Points, Scale, Score, Translation, VisionFrame, WarpTemplate, WarpTemplateSet

WARP_TEMPLATE_SET : WarpTemplateSet =\
{
	'arcface_112_v1': numpy.array(
	[
		[ 0.35473214, 0.45658929 ],
		[ 0.64526786, 0.45658929 ],
		[ 0.50000000, 0.61154464 ],
		[ 0.37913393, 0.77687500 ],
		[ 0.62086607, 0.77687500 ]
	]),
	'arcface_112_v2': numpy.array(
	[
		[ 0.34191607, 0.46157411 ],
		[ 0.65653393, 0.45983393 ],
		[ 0.50022500, 0.64050536 ],
		[ 0.37097589, 0.82469196 ],
		[ 0.63151696, 0.82325089 ]
	]),
	'arcface_128': numpy.array(
	[
		[ 0.36167656, 0.40387734 ],
		[ 0.63696719, 0.40235469 ],
		[ 0.50019687, 0.56044219 ],
		[ 0.38710391, 0.72160547 ],
		[ 0.61507734, 0.72034453 ]
	]),
	'dfl_whole_face': numpy.array(
	[
		[ 0.35342266, 0.39285716 ],
		[ 0.62797622, 0.39285716 ],
		[ 0.48660713, 0.54017860 ],
		[ 0.38839287, 0.68750011 ],
		[ 0.59821427, 0.68750011 ]
	]),
	'ffhq_512': numpy.array(
	[
		[ 0.37691676, 0.46864664 ],
		[ 0.62285697, 0.46912813 ],
		[ 0.50123859, 0.61331904 ],
		[ 0.39308822, 0.72541100 ],
		[ 0.61150205, 0.72490465 ]
	]),
	'mtcnn_512': numpy.array(
	[
		[ 0.36562865, 0.46733799 ],
		[ 0.63305391, 0.46585885 ],
		[ 0.50019127, 0.61942959 ],
		[ 0.39032951, 0.77598822 ],
		[ 0.61178945, 0.77476328 ]
	]),
	'styleganex_384': numpy.array(
	[
		[ 0.42353745, 0.52289879 ],
		[ 0.57725008, 0.52319972 ],
		[ 0.50123859, 0.61331904 ],
		[ 0.43364461, 0.68337652 ],
		[ 0.57015325, 0.68306005 ]
	])
}


def estimate_matrix_by_face_landmark_5(face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Matrix:
	warp_template_norm = WARP_TEMPLATE_SET.get(warp_template) * crop_size
	affine_matrix = cv2.estimateAffinePartial2D(face_landmark_5, warp_template_norm, method = cv2.RANSAC, ransacReprojThreshold = 100)[0]
	return affine_matrix


def warp_face_by_face_landmark_5(temp_vision_frame : VisionFrame, face_landmark_5 : FaceLandmark5, warp_template : WarpTemplate, crop_size : Size) -> Tuple[VisionFrame, Matrix]:
	affine_matrix = estimate_matrix_by_face_landmark_5(face_landmark_5, warp_template, crop_size)
	crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size, borderMode = cv2.BORDER_REPLICATE, flags = cv2.INTER_AREA)
	return crop_vision_frame, affine_matrix


def warp_face_by_bounding_box(temp_vision_frame : VisionFrame, bounding_box : BoundingBox, crop_size : Size) -> Tuple[VisionFrame, Matrix]:
	source_points = numpy.array([ [ bounding_box[0], bounding_box[1] ], [ bounding_box[2], bounding_box[1] ], [ bounding_box[0], bounding_box[3] ] ]).astype(numpy.float32)
	target_points = numpy.array([ [ 0, 0 ], [ crop_size[0], 0 ], [ 0, crop_size[1] ] ]).astype(numpy.float32)
	affine_matrix = cv2.getAffineTransform(source_points, target_points)
	if bounding_box[2] - bounding_box[0] > crop_size[0] or bounding_box[3] - bounding_box[1] > crop_size[1]:
		interpolation_method = cv2.INTER_AREA
	else:
		interpolation_method = cv2.INTER_LINEAR
	crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size, flags = interpolation_method)
	return crop_vision_frame, affine_matrix


def warp_face_by_translation(temp_vision_frame : VisionFrame, translation : Translation, scale : float, crop_size : Size) -> Tuple[VisionFrame, Matrix]:
	affine_matrix = numpy.array([ [ scale, 0, translation[0] ], [ 0, scale, translation[1] ] ])
	crop_vision_frame = cv2.warpAffine(temp_vision_frame, affine_matrix, crop_size)
	return crop_vision_frame, affine_matrix


def paste_back(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, crop_vision_mask : Mask, affine_matrix : Matrix) -> VisionFrame:
	paste_bounding_box, paste_matrix = calculate_paste_area(temp_vision_frame, crop_vision_frame, affine_matrix)
	x1, y1, x2, y2 = paste_bounding_box
	paste_width = x2 - x1
	paste_height = y2 - y1
	inverse_vision_mask = cv2.warpAffine(crop_vision_mask, paste_matrix, (paste_width, paste_height)).clip(0, 1)
	inverse_vision_mask = numpy.expand_dims(inverse_vision_mask, axis = -1)
	inverse_vision_frame = cv2.warpAffine(crop_vision_frame, paste_matrix, (paste_width, paste_height), borderMode = cv2.BORDER_REPLICATE)
	temp_vision_frame = temp_vision_frame.copy()
	paste_vision_frame = temp_vision_frame[y1:y2, x1:x2]
	paste_vision_frame = paste_vision_frame * (1 - inverse_vision_mask) + inverse_vision_frame * inverse_vision_mask
	temp_vision_frame[y1:y2, x1:x2] = paste_vision_frame.astype(temp_vision_frame.dtype)
	return temp_vision_frame


def calculate_paste_area(temp_vision_frame : VisionFrame, crop_vision_frame : VisionFrame, affine_matrix : Matrix) -> Tuple[BoundingBox, Matrix]:
	temp_height, temp_width = temp_vision_frame.shape[:2]
	crop_height, crop_width = crop_vision_frame.shape[:2]
	inverse_matrix = cv2.invertAffineTransform(affine_matrix)
	crop_points = numpy.array([ [ 0, 0 ], [ crop_width, 0 ], [ crop_width, crop_height ], [ 0, crop_height ] ])
	paste_region_points = transform_points(crop_points, inverse_matrix)
	paste_region_point_min = numpy.floor(paste_region_points.min(axis = 0)).astype(int)
	paste_region_point_max = numpy.ceil(paste_region_points.max(axis = 0)).astype(int)
	x1, y1 = numpy.clip(paste_region_point_min, 0, [ temp_width, temp_height ])
	x2, y2 = numpy.clip(paste_region_point_max, 0, [ temp_width, temp_height ])
	paste_bounding_box = numpy.array([ x1, y1, x2, y2 ])
	paste_matrix = inverse_matrix.copy()
	paste_matrix[0, 2] -= x1
	paste_matrix[1, 2] -= y1
	return paste_bounding_box, paste_matrix


@lru_cache()
def create_static_anchors(feature_stride : int, anchor_total : int, stride_height : int, stride_width : int) -> Anchors:
	x, y = numpy.mgrid[:stride_width, :stride_height]
	anchors = numpy.stack((y, x), axis = -1)
	anchors = (anchors * feature_stride).reshape((-1, 2))
	anchors = numpy.stack([ anchors ] * anchor_total, axis = 1).reshape((-1, 2))
	return anchors


def create_rotation_matrix_and_size(angle : Angle, size : Size) -> Tuple[Matrix, Size]:
	rotation_matrix = cv2.getRotationMatrix2D((size[0] / 2, size[1] / 2), angle, 1)
	rotation_size = numpy.dot(numpy.abs(rotation_matrix[:, :2]), size)
	rotation_matrix[:, -1] += (rotation_size - size) * 0.5 #type:ignore[misc]
	rotation_size = int(rotation_size[0]), int(rotation_size[1])
	return rotation_matrix, rotation_size


def create_bounding_box(face_landmark_68 : FaceLandmark68) -> BoundingBox:
	x1, y1 = numpy.min(face_landmark_68, axis = 0)
	x2, y2 = numpy.max(face_landmark_68, axis = 0)
	bounding_box = normalize_bounding_box(numpy.array([ x1, y1, x2, y2 ]))
	return bounding_box


def normalize_bounding_box(bounding_box : BoundingBox) -> BoundingBox:
	x1, y1, x2, y2 = bounding_box
	x1, x2 = sorted([ x1, x2 ])
	y1, y2 = sorted([ y1, y2 ])
	return numpy.array([ x1, y1, x2, y2 ])


def transform_points(points : Points, matrix : Matrix) -> Points:
	points = points.reshape(-1, 1, 2)
	points = cv2.transform(points, matrix) #type:ignore[assignment]
	points = points.reshape(-1, 2)
	return points


def transform_bounding_box(bounding_box : BoundingBox, matrix : Matrix) -> BoundingBox:
	points = numpy.array(
	[
		[ bounding_box[0], bounding_box[1] ],
		[ bounding_box[2], bounding_box[1] ],
		[ bounding_box[2], bounding_box[3] ],
		[ bounding_box[0], bounding_box[3] ]
	])
	points = transform_points(points, matrix)
	x1, y1 = numpy.min(points, axis = 0)
	x2, y2 = numpy.max(points, axis = 0)
	return normalize_bounding_box(numpy.array([ x1, y1, x2, y2 ]))


def distance_to_bounding_box(points : Points, distance : Distance) -> BoundingBox:
	x1 = points[:, 0] - distance[:, 0]
	y1 = points[:, 1] - distance[:, 1]
	x2 = points[:, 0] + distance[:, 2]
	y2 = points[:, 1] + distance[:, 3]
	bounding_box = numpy.column_stack([ x1, y1, x2, y2 ])
	return bounding_box


def distance_to_face_landmark_5(points : Points, distance : Distance) -> FaceLandmark5:
	x = points[:, 0::2] + distance[:, 0::2]
	y = points[:, 1::2] + distance[:, 1::2]
	face_landmark_5 = numpy.stack((x, y), axis = -1)
	return face_landmark_5


def scale_face_landmark_5(face_landmark_5 : FaceLandmark5, scale : Scale) -> FaceLandmark5:
	face_landmark_5_scale = face_landmark_5 - face_landmark_5[2]
	face_landmark_5_scale *= scale
	face_landmark_5_scale += face_landmark_5[2]
	return face_landmark_5_scale


def convert_to_face_landmark_5(face_landmark_68 : FaceLandmark68) -> FaceLandmark5:
	face_landmark_5 = numpy.array(
	[
		numpy.mean(face_landmark_68[36:42], axis = 0),
		numpy.mean(face_landmark_68[42:48], axis = 0),
		face_landmark_68[30],
		face_landmark_68[48],
		face_landmark_68[54]
	])
	return face_landmark_5


def estimate_face_angle(face_landmark_68 : FaceLandmark68) -> Angle:
	x1, y1 = face_landmark_68[0]
	x2, y2 = face_landmark_68[16]
	theta = numpy.arctan2(y2 - y1, x2 - x1)
	theta = numpy.degrees(theta) % 360
	angles = numpy.linspace(0, 360, 5)
	index = numpy.argmin(numpy.abs(angles - theta))
	face_angle = int(angles[index] % 360)
	return face_angle


def apply_nms(bounding_boxes : List[BoundingBox], scores : List[Score], score_threshold : float, nms_threshold : float) -> Sequence[int]:
	bounding_boxes_norm = [ (x1, y1, x2 - x1, y2 - y1) for (x1, y1, x2, y2) in bounding_boxes ]
	keep_indices = cv2.dnn.NMSBoxes(bounding_boxes_norm, scores, score_threshold = score_threshold, nms_threshold = nms_threshold)
	return keep_indices


def get_nms_threshold(face_detector_model : FaceDetectorModel, face_detector_angles : List[Angle]) -> float:
	if face_detector_model == 'many':
		return 0.1
	if len(face_detector_angles) == 2:
		return 0.3
	if len(face_detector_angles) == 3:
		return 0.2
	if len(face_detector_angles) == 4:
		return 0.1
	return 0.4


def merge_matrix(temp_matrices : List[Matrix]) -> Matrix:
	matrix = numpy.vstack([ temp_matrices[0], [ 0, 0, 1 ] ])

	for temp_matrix in temp_matrices[1:]:
		temp_matrix = numpy.vstack([ temp_matrix, [ 0, 0, 1 ] ])
		matrix = numpy.dot(temp_matrix, matrix)

	return matrix[:2, :]


def calculate_bounding_box_overlap(bounding_box_a : BoundingBox, bounding_box_b : BoundingBox) -> float:
	intersection_x1 = max(bounding_box_a[0], bounding_box_b[0])
	intersection_y1 = max(bounding_box_a[1], bounding_box_b[1])
	intersection_x2 = min(bounding_box_a[2], bounding_box_b[2])
	intersection_y2 = min(bounding_box_a[3], bounding_box_b[3])
	intersection = max(0, intersection_x2 - intersection_x1) * max(0, intersection_y2 - intersection_y1)
	bounding_box_area = (bounding_box_a[2] - bounding_box_a[0]) * (bounding_box_a[3] - bounding_box_a[1])
	reference_bounding_box_area = (bounding_box_b[2] - bounding_box_b[0]) * (bounding_box_b[3] - bounding_box_b[1])
	union = bounding_box_area + reference_bounding_box_area - intersection

	if union > 0:
		return intersection / union

	return 0.0


def average_points(points_previous : Points, points_next : Points, average_factor : float) -> Points:
	return points_previous * (1 - average_factor) + points_next * average_factor