import torch
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm

from configs.config import cfg
from .ProjectionUtils import HomoCalc
from .ProjectionUtils import HomoProj


def mesh_warp_frame(frame, x_motion, y_motion):
    """
    @param frame current frame [N, 1, H, W]
    @param x_motion [N, 1, G_H, G_W]
    @param y_motion [N, 1, G_H, G_W]

    @return mesh warping according to given motion
    """

    target_device = frame.device

    src_grids = torch.stack(torch.meshgrid(torch.arange(0, cfg.MODEL.WIDTH, cfg.MODEL.PIXELS, device=target_device),
                                           torch.arange(0, cfg.MODEL.HEIGHT, cfg.MODEL.PIXELS, device=target_device)), 0).permute(0, 2, 1).unsqueeze(0).float()  # 2, G_H, G_W

    des_grids = src_grids + torch.cat([x_motion, y_motion], 1)

    projection = []

    for i in range(des_grids.shape[0]):
        homo = HomoCalc(src_grids[0], des_grids[i])

        origin_kp = torch.stack(torch.meshgrid(torch.arange(0, cfg.MODEL.WIDTH, device=target_device),
                                               torch.arange(0, cfg.MODEL.HEIGHT, device=target_device)), 0).permute(0, 2, 1).float()  # 2, H, W

        projected_kp = HomoProj(homo, origin_kp.contiguous().view(
            2, -1).permute(1, 0)).permute(1, 0)

        projection.append(projected_kp.contiguous().view(
            *origin_kp.shape).permute(1, 2, 0))  # 2, H, W --> H, W, 2
    projection = torch.stack(projection, 0)

    projection[:, :, :, 0] = projection[:, :, :, 0] / cfg.MODEL.WIDTH * 2. - 1.
    projection[:, :, :, 1] = projection[:, :, :, 1] / \
        cfg.MODEL.HEIGHT * 2. - 1.
    generated_frame = F.grid_sample(frame, projection, align_corners=True)

    return generated_frame


def warpListImage(images, x_motion, y_motion):
    """
    @param images List(image [1, 1, H, W]) 
    @param x_motion [G_H, G_W, N]
    @param y_motion [G_H, G_W, N]
    """

    frames = np.concatenate(images, 0)
    x_motion = np.expand_dims(np.transpose(x_motion, (2, 0, 1)), 1)
    y_motion = np.expand_dims(np.transpose(y_motion, (2, 0, 1)), 1)
    frames = torch.from_numpy(frames.astype(np.float32))
    x_motion = torch.from_numpy(x_motion.astype(np.float32))
    y_motion = torch.from_numpy(y_motion.astype(np.float32))
    return mesh_warp_frame(frames, x_motion, y_motion)
