from typing import List, Dict, Callable, Tuple, Optional
import torch
import torch.nn.functional as F
import functools
import numpy as np


def get_crop_and_resize_matrix(
        box: torch.Tensor, target_shape: Tuple[int, int],
        target_face_scale: float = 1.0, make_square_crop: bool = True,
        offset_xy: Optional[Tuple[float, float]] = None, align_corners: bool = True,
        offset_box_coords: bool = False) -> torch.Tensor:
    """
    Args:
        box: b x 4(x1, y1, x2, y2)
        align_corners (bool): Set this to `True` only if the box you give has coordinates
            ranging from `0` to `h-1` or `w-1`.

        offset_box_coords (bool): Set this to `True` if the box you give has coordinates
            ranging from `0` to `h` or `w`. 

            Set this to `False` if the box coordinates range from `-0.5` to `h-0.5` or `w-0.5`.

            If the box coordinates range from `0` to `h-1` or `w-1`, set `align_corners=True`.

    Returns:
        torch.Tensor: b x 3 x 3.
    """
    if offset_xy is None:
        offset_xy = (0.0, 0.0)

    x1, y1, x2, y2 = box.split(1, dim=1)  # b x 1
    cx = (x1 + x2) / 2 + offset_xy[0]
    cy = (y1 + y2) / 2 + offset_xy[1]
    rx = (x2 - x1) / 2 / target_face_scale
    ry = (y2 - y1) / 2 / target_face_scale
    if make_square_crop:
        rx = ry = torch.maximum(rx, ry)

    x1, y1, x2, y2 = cx - rx, cy - ry, cx + rx, cy + ry

    h, w, *_ = target_shape

    zeros_pl = torch.zeros_like(x1)
    ones_pl = torch.ones_like(x1)

    if align_corners:
        # x -> (x - x1) / (x2 - x1) * (w - 1)
        # y -> (y - y1) / (y2 - y1) * (h - 1)
        ax = 1.0 / (x2 - x1) * (w - 1)
        ay = 1.0 / (y2 - y1) * (h - 1)
        matrix = torch.cat([
            ax, zeros_pl, -x1 * ax,
            zeros_pl, ay, -y1 * ay,
            zeros_pl, zeros_pl, ones_pl
        ], dim=1).reshape(-1, 3, 3)  # b x 3 x 3
    else:
        if offset_box_coords:
            # x1, x2 \in [0, w], y1, y2 \in [0, h]
            # first we should offset x1, x2, y1, y2 to be ranging in
            # [-0.5, w-0.5] and [-0.5, h-0.5]
            # so to convert these pixel coordinates into boundary coordinates.
            x1, x2, y1, y2 = x1-0.5, x2-0.5, y1-0.5, y2-0.5

        # x -> (x - x1) / (x2 - x1) * w - 0.5
        # y -> (y - y1) / (y2 - y1) * h - 0.5
        ax = 1.0 / (x2 - x1) * w
        ay = 1.0 / (y2 - y1) * h
        matrix = torch.cat([
            ax, zeros_pl, -x1 * ax - 0.5*ones_pl,
            zeros_pl, ay, -y1 * ay - 0.5*ones_pl,
            zeros_pl, zeros_pl, ones_pl
        ], dim=1).reshape(-1, 3, 3)  # b x 3 x 3
    return matrix


def get_similarity_transform_matrix(
        from_pts: torch.Tensor, to_pts: torch.Tensor) -> torch.Tensor:
    """
    Args:
        from_pts, to_pts: b x n x 2

    Returns:
        torch.Tensor: b x 3 x 3
    """
    mfrom = from_pts.mean(dim=1, keepdim=True)  # b x 1 x 2
    mto = to_pts.mean(dim=1, keepdim=True)  # b x 1 x 2

    a1 = (from_pts - mfrom).square().sum([1, 2], keepdim=False)  # b
    c1 = ((to_pts - mto) * (from_pts - mfrom)).sum([1, 2], keepdim=False)  # b

    to_delta = to_pts - mto
    from_delta = from_pts - mfrom
    c2 = (to_delta[:, :, 0] * from_delta[:, :, 1] - to_delta[:,
          :, 1] * from_delta[:, :, 0]).sum([1], keepdim=False)  # b

    a = c1 / a1
    b = c2 / a1
    dx = mto[:, 0, 0] - a * mfrom[:, 0, 0] - b * mfrom[:, 0, 1]  # b
    dy = mto[:, 0, 1] + b * mfrom[:, 0, 0] - a * mfrom[:, 0, 1]  # b

    ones_pl = torch.ones_like(a1)
    zeros_pl = torch.zeros_like(a1)

    return torch.stack([
        a, b, dx,
        -b, a, dy,
        zeros_pl, zeros_pl, ones_pl,
    ], dim=-1).reshape(-1, 3, 3)


@functools.lru_cache()
def _standard_face_pts():
    pts = torch.tensor([
        196.0, 226.0,
        316.0, 226.0,
        256.0, 286.0,
        220.0, 360.4,
        292.0, 360.4], dtype=torch.float32) / 256.0 - 1.0
    return torch.reshape(pts, (5, 2))


def get_face_align_matrix(
        face_pts: torch.Tensor, target_shape: Tuple[int, int],
        target_face_scale: float = 1.0, offset_xy: Optional[Tuple[float, float]] = None,
        target_pts: Optional[torch.Tensor] = None):

    if target_pts is None:
        with torch.no_grad():
            std_pts = _standard_face_pts().to(face_pts)  # [-1 1]
            h, w, *_ = target_shape
            target_pts = (std_pts * target_face_scale + 1) * \
                torch.tensor([w-1, h-1]).to(face_pts) / 2.0
            if offset_xy is not None:
                target_pts[:, 0] += offset_xy[0]
                target_pts[:, 1] += offset_xy[1]
    else:
        target_pts = target_pts.to(face_pts)

    if target_pts.dim() == 2:
        target_pts = target_pts.unsqueeze(0)
    if target_pts.size(0) == 1:
        target_pts = target_pts.broadcast_to(face_pts.shape)

    assert target_pts.shape == face_pts.shape

    return get_similarity_transform_matrix(face_pts, target_pts)


def rot90(v):
    return np.array([-v[1], v[0]])


def get_quad(lm: torch.Tensor):
    # N,2
    lm = lm.detach().cpu().numpy()
    # Choose oriented crop rectangle.
    eye_avg = (lm[0] + lm[1]) * 0.5 + 0.5
    mouth_avg = (lm[3] + lm[4]) * 0.5 + 0.5
    eye_to_eye = lm[1] - lm[0]
    eye_to_mouth = mouth_avg - eye_avg
    x = eye_to_eye - rot90(eye_to_mouth)
    x /= np.hypot(*x)
    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
    y = rot90(x)
    c = eye_avg + eye_to_mouth * 0.1
    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
    quad_for_coeffs = quad[[0,3, 2,1]] #  顺序改一下
    return torch.from_numpy(quad_for_coeffs).float()


def get_face_align_matrix_celebm(
        face_pts: torch.Tensor, target_shape: Tuple[int, int]):

    face_pts = torch.stack([get_quad(pts) for pts in face_pts], dim=0).to(face_pts)

    assert target_shape[0] == target_shape[1]
    target_size  = target_shape[0]
    target_pts = torch.as_tensor([[0, 0], [target_size,0], [target_size, target_size], [0, target_size]]).to(face_pts)

    if target_pts.dim() == 2:
        target_pts = target_pts.unsqueeze(0)
    if target_pts.size(0) == 1:
        target_pts = target_pts.broadcast_to(face_pts.shape)

    assert target_pts.shape == face_pts.shape

    return get_similarity_transform_matrix(face_pts, target_pts)

@functools.lru_cache(maxsize=128)
def _meshgrid(h, w) -> Tuple[torch.Tensor, torch.Tensor]:
    yy, xx = torch.meshgrid(torch.arange(h).float(),
                            torch.arange(w).float(),
                            indexing='ij')
    return yy, xx


def _forge_grid(batch_size: int, device: torch.device,
                output_shape: Tuple[int, int],
                fn: Callable[[torch.Tensor], torch.Tensor]
                ) -> Tuple[torch.Tensor, torch.Tensor]:
    """ Forge transform maps with a given function `fn`.

    Args:
        output_shape (tuple): (b, h, w, ...).
        fn (Callable[[torch.Tensor], torch.Tensor]): The function that accepts 
            a bxnx2 array and outputs the transformed bxnx2 array. Both input 
            and output store (x, y) coordinates.

    Note: 
        both input and output arrays of `fn` should store (y, x) coordinates.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: Two maps `X` and `Y`, where for each 
            pixel (y, x) or coordinate (x, y),
            `(X[y, x], Y[y, x]) = fn([x, y])`
    """
    h, w, *_ = output_shape
    yy, xx = _meshgrid(h, w)  # h x w
    yy = yy.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)
    xx = xx.unsqueeze(0).broadcast_to(batch_size, h, w).to(device)

    in_xxyy = torch.stack(
        [xx, yy], dim=-1).reshape([batch_size, h*w, 2])  # (h x w) x 2
    out_xxyy: torch.Tensor = fn(in_xxyy)  # (h x w) x 2
    return out_xxyy.reshape(batch_size, h, w, 2)


def _safe_arctanh(x: torch.Tensor, eps: float = 0.001) -> torch.Tensor:
    return torch.clamp(x, -1+eps, 1-eps).arctanh()


def inverted_tanh_warp_transform(coords: torch.Tensor, matrix: torch.Tensor,
                                 warp_factor: float, warped_shape: Tuple[int, int]):
    """ Inverted tanh-warp function.

    Args:
        coords (torch.Tensor): b x n x 2 (x, y). The transformed coordinates.
        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
            from the original image to the aligned yet not-warped image.
        warp_factor (float): The warp factor. 
            0 means linear transform, 1 means full tanh warp.
        warped_shape (tuple): [height, width].

    Returns:
        torch.Tensor: b x n x 2 (x, y). The original coordinates.
    """
    h, w, *_ = warped_shape
    # h -= 1
    # w -= 1

    w_h = torch.tensor([[w, h]]).to(coords)

    if warp_factor > 0:
        # normalize coordinates to [-1, +1]
        coords = coords / w_h * 2 - 1

        nl_part1 = coords > 1.0 - warp_factor
        nl_part2 = coords < -1.0 + warp_factor

        ret_nl_part1 = _safe_arctanh(
            (coords - 1.0 + warp_factor) /
            warp_factor) * warp_factor + \
            1.0 - warp_factor
        ret_nl_part2 = _safe_arctanh(
            (coords + 1.0 - warp_factor) /
            warp_factor) * warp_factor - \
            1.0 + warp_factor

        coords = torch.where(nl_part1, ret_nl_part1,
                             torch.where(nl_part2, ret_nl_part2, coords))

        # denormalize
        coords = (coords + 1) / 2 * w_h

    coords_homo = torch.cat(
        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3

    # inv_matrix = torch.linalg.inv(matrix)  # b x 3 x 3
    device = matrix.device
    inv_matrix_np = np.linalg.inv(matrix.cpu().numpy())
    inv_matrix = torch.from_numpy(inv_matrix_np).to(device)
    coords_homo = torch.bmm(
        coords_homo, inv_matrix.permute(0, 2, 1))  # b x n x 3
    return coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]]


def tanh_warp_transform(
        coords: torch.Tensor, matrix: torch.Tensor,
        warp_factor: float, warped_shape: Tuple[int, int]):
    """ Tanh-warp function.

    Args:
        coords (torch.Tensor): b x n x 2 (x, y). The original coordinates.
        matrix: b x 3 x 3. A matrix that transforms un-normalized coordinates 
            from the original image to the aligned yet not-warped image.
        warp_factor (float): The warp factor. 
            0 means linear transform, 1 means full tanh warp.
        warped_shape (tuple): [height, width].

    Returns:
        torch.Tensor: b x n x 2 (x, y). The transformed coordinates.
    """
    h, w, *_ = warped_shape
    # h -= 1
    # w -= 1
    w_h = torch.tensor([[w, h]]).to(coords)

    coords_homo = torch.cat(
        [coords, torch.ones_like(coords[:, :, [0]])], dim=-1)  # b x n x 3

    coords_homo = torch.bmm(coords_homo, matrix.transpose(2, 1))  # b x n x 3
    coords = (coords_homo[:, :, :2] / coords_homo[:, :, [2, 2]])  # b x n x 2

    if warp_factor > 0:
        # normalize coordinates to [-1, +1]
        coords = coords / w_h * 2 - 1

        nl_part1 = coords > 1.0 - warp_factor
        nl_part2 = coords < -1.0 + warp_factor

        ret_nl_part1 = torch.tanh(
            (coords - 1.0 + warp_factor) /
            warp_factor) * warp_factor + \
            1.0 - warp_factor
        ret_nl_part2 = torch.tanh(
            (coords + 1.0 - warp_factor) /
            warp_factor) * warp_factor - \
            1.0 + warp_factor

        coords = torch.where(nl_part1, ret_nl_part1,
                             torch.where(nl_part2, ret_nl_part2, coords))

        # denormalize
        coords = (coords + 1) / 2 * w_h

    return coords


def make_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
                        warped_shape: Tuple[int, int],
                        orig_shape: Tuple[int, int]):
    """
    Args:
        matrix: bx3x3 matrix.
        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
           `warp_factor=0.0` represents a cropping.
        warped_shape: The target image shape to transform to.

    Returns:
        torch.Tensor: b x h x w x 2 (x, y).
    """
    orig_h, orig_w, *_ = orig_shape
    w_h = torch.tensor([orig_w, orig_h]).to(matrix).reshape(1, 1, 1, 2)
    return _forge_grid(
        matrix.size(0), matrix.device,
        warped_shape,
        functools.partial(inverted_tanh_warp_transform,
                          matrix=matrix,
                          warp_factor=warp_factor,
                          warped_shape=warped_shape)) / w_h*2-1


def make_inverted_tanh_warp_grid(matrix: torch.Tensor, warp_factor: float,
                                 warped_shape: Tuple[int, int],
                                 orig_shape: Tuple[int, int]):
    """
    Args:
        matrix: bx3x3 matrix.
        warp_factor: The warping factor. `warp_factor=1.0` represents a vannila Tanh-warping, 
           `warp_factor=0.0` represents a cropping.
        warped_shape: The target image shape to transform to.
        orig_shape: The original image shape that is transformed from.

    Returns:
        torch.Tensor: b x h x w x 2 (x, y).
    """
    h, w, *_ = warped_shape
    w_h = torch.tensor([w, h]).to(matrix).reshape(1, 1, 1, 2)
    return _forge_grid(
        matrix.size(0), matrix.device,
        orig_shape,
        functools.partial(tanh_warp_transform,
                          matrix=matrix,
                          warp_factor=warp_factor,
                          warped_shape=warped_shape)) / w_h * 2-1