from typing import Optional, Tuple, List, Union
import numpy as np
from math import e, pi

from model_demo.inference.constants import BLENDSHAPE_NAMES


def clean_up_blendshapes(
    blendshapes: np.ndarray,
    mouth_exaggeration: float,
    brow_exaggeration: float,
    unsquinch_fix: float,
    eye_contact_fix: float,
    clear_neutral: bool = False,
    exaggerate_above: float = 0,
    symmetrize_eyes: bool = False,
) -> np.ndarray:
    """
    Exaggerate blendshapes by a given factor.

    Args:
        blendshapes: Blendshape coefficients of shape (B, T, D) or (T, D)
        exaggeration_factor: Factor to exaggerate the blendshapes by
        unsquinch_fix: Factor to reduce eye squint and blink blendshapes by in range [0, 1]
        eye_contact_fix: Factor to reduce eye look blendshapes by in range [0, 1]
        clear_neutral: Whether to clear the neutral expression blendshape (set to 0)
        mouth_only: Whether to exaggerate only the mouth blendshapes
        exaggerate_above: Landmarks below this value will be exaggerated up, below down

    Returns:
        Exaggerated blendshape coefficients of shape (B, T, D) or (T, D)
    """

    def modify_blendshapes(
        blendshapes: np.ndarray, target_substrings: List[str], factor: float
    ) -> np.ndarray:
        if factor != 1:
            for i, shape in enumerate(BLENDSHAPE_NAMES):
                if any(substring in shape for substring in target_substrings):
                    blendshapes_offset = blendshapes[..., i] - exaggerate_above
                    blendshapes[..., i] = blendshapes_offset * factor + exaggerate_above
        blendshapes = np.clip(blendshapes, 0.0, 1.0)
        return blendshapes

    if clear_neutral:
        blendshapes[..., 0] = 0

    modify_blendshapes(blendshapes, ["mouth", "jaw", "cheek"], mouth_exaggeration)
    modify_blendshapes(blendshapes, ["brow", "noseSneer", "eye"], brow_exaggeration)
    if unsquinch_fix > 0:
        eye_idx = [
            i
            for i, name in enumerate(BLENDSHAPE_NAMES)
            if "eyeSquint" in name or "eyeBlink" in name
        ]
        for idx in eye_idx:
            blendshapes[..., idx] -= unsquinch_fix
    if eye_contact_fix > 0:
        eye_idx = [i for i, name in enumerate(BLENDSHAPE_NAMES) if "eyeLook" in name]
        for idx in eye_idx:
            blendshapes[..., idx] -= eye_contact_fix
    if symmetrize_eyes:
        # average between eyeBlinkLeft and eyeBlinkRight
        eye_blink_left_index = BLENDSHAPE_NAMES.index("eyeBlinkLeft")
        eye_blink_right_index = BLENDSHAPE_NAMES.index("eyeBlinkRight")
        avg_val = (
            blendshapes[..., eye_blink_left_index]
            + blendshapes[..., eye_blink_right_index]
        ) / 2
        blendshapes[..., eye_blink_left_index] = avg_val
        blendshapes[..., eye_blink_right_index] = avg_val

    blendshapes = np.clip(blendshapes, 0.0, 1.0)

    return blendshapes


def exaggerate_head_wiggle(
    head_angles: np.ndarray[np.float32], exaggeration_factor: float
) -> np.ndarray[np.float32]:
    """
    Exaggerate head angles by a given factor.

    Args:
        head_angles: Sequence of pitch, yaw, roll values of shape (temporal_dim, 3)
        exaggeration_factor: Factor to exaggerate the head angles by

    Returns:
        Exaggerated head angles of shape (temporal_dim, 3)
    """
    return head_angles * exaggeration_factor


def unscale_and_uncenter_head_angles(
    head_angles: np.ndarray[np.float32],
    mean_pos: Optional[np.ndarray[np.float32]] = None,
    bad_frames: List[int] = [],
) -> np.ndarray[np.float32]:
    """
    Rescale head angles in range [-1, 1] to [-pi, pi] and uncenter them.

    Args:
        head_angles: Sequence of pitch, yaw, roll values of shape (temporal_dim, 3)
        mean_pos: Mean position to offset the head angles of shape (3,)
        bad_frames: List of indices of frames where face detection failed

    Returns:
        Array of unscaled and uncentered head angles of shape (temporal_dim, 3)
    """
    if mean_pos is None:
        mean_pos = np.zeros(3).astype(np.float32)
    good_frames = [i for i in range(head_angles.shape[0]) if i not in bad_frames]
    head_angles[good_frames] = head_angles[good_frames] + mean_pos
    head_angles[good_frames] = head_angles[good_frames] * pi
    return head_angles