#! /usr/bin/env python
# -*- coding: utf-8 -*-

"""
Image deformation using moving least squares

@author: <NAME>
@email: <EMAIL>
@date: 2017/8/8
@update: 2020/9/25
@update: 2021/7/14: Simplify usage
@update: 2021/12/24: Fix bugs and add an example of random control points (see `demo2()`)
"""

import time
import cv2
import torch
import numpy as np
import matplotlib.pyplot as plt

try:
    import torch    # Install PyTorch first: https://pytorch.org/get-started/locally/
    from img_utils_pytorch import (
        mls_affine_deformation as mls_affine_deformation_pt,
        mls_similarity_deformation as mls_similarity_deformation_pt,
        mls_rigid_deformation as mls_rigid_deformation_pt,
    )
    device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
except ImportError as e:
    print(e)

from img_utils import (
    mls_affine_deformation, 
    mls_similarity_deformation, 
    mls_rigid_deformation
)

from PIL import Image




def MLS(image, input_points, output_points):

    height, width, _ = image.shape
    gridX = np.arange(width, dtype=np.int16)
    gridY = np.arange(height, dtype=np.int16)
    vx, vy = np.meshgrid(gridX, gridY)
    
    
    start = time.time()
    #仿射变形,速度最快
    affine = mls_affine_deformation(vy, vx, input_points, output_points, alpha=1)
    UX =  affine[0,:,:].astype(np.float32)
    UY =  affine[1,:,:].astype(np.float32)
    aug1 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    
    end1 = time.time()
    print("MLS time cost:{} s".format(end1 - start))
    
    
    #相似变形,速度最慢
    similar = mls_similarity_deformation(vy, vx, input_points, output_points, alpha=1)
    UX =  affine[0,:,:].astype(np.float32)
    UY =  affine[1,:,:].astype(np.float32)
    aug2 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    
    end2 = time.time()
    print("MLS time cost:{} s".format(end2 - end1))


    #网格变形,速度中间和最慢差不多
    rigid = mls_rigid_deformation(vy, vx, input_points, output_points, alpha=1)
    UX =  affine[0,:,:].astype(np.float32)
    UY =  affine[1,:,:].astype(np.float32)
    aug3 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)

    end3 = time.time()
    print("MLS time cost:{} s".format(end3 - end2))

   
    
    end = time.time()
    print("MLS time cost:{} s".format(end - start))
    return aug1











def MLS_TORCH(image, input_points, output_points):

    device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
    input_points_pytorch = torch.from_numpy(input_points).to(device)
    output_points_pytorch = torch.from_numpy(output_points).to(device)
    image_pytorch = torch.from_numpy(image).to(device)


    height, width, _ = image_pytorch.shape
    gridX = torch.arange(width, dtype=torch.int16).to(device)
    gridY = torch.arange(height, dtype=torch.int16).to(device)
    vx, vy = torch.meshgrid(gridX, gridY)
    # !!! Pay attention !!!: the shape of returned tensors are different between numpy.meshgrid and torch.meshgrid
    vx, vy = vx.transpose(0, 1), vy.transpose(0, 1)


    
    start = time.time()
    #仿射变形
    affine = mls_affine_deformation_pt(vy, vx, input_points_pytorch, output_points_pytorch, alpha=1)
    affine = affine.cpu().numpy()
    UX =  affine[0,:,:].astype(np.float32)
    UY =  affine[1,:,:].astype(np.float32)
    aug1 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    
    end1 = time.time()
    print("MLS TORCH time cost:{} s".format(end1 - start))
 

    
    #相似变形
    similar = mls_similarity_deformation_pt(vy, vx, input_points_pytorch, output_points_pytorch, alpha=1)
    similar = similar.cpu().numpy()
    UX =  similar[0,:,:].astype(np.float32)
    UY =  similar[1,:,:].astype(np.float32)
    aug2 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    
    end2 = time.time()
    print("MLS TORCH time cost:{} s".format(end2 - end1))
 



    #网格变形
    rigid = mls_rigid_deformation_pt(vy, vx, input_points_pytorch, output_points_pytorch, alpha=1)
    rigid = rigid.cpu().numpy()
    UX =  rigid[0,:,:].astype(np.float32)
    UY =  rigid[1,:,:].astype(np.float32)
    aug3 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    
    end3 = time.time()
    print("MLS TORCH time cost:{} s".format(end3 - end2))
 

    end = time.time()
    print("MLS TORCH time cost:{} s".format(end - start))
    return aug1




def demo():
    p = np.array([
        [155, 30], [155, 125], [155, 225],
        [235, 100], [235, 160], [295, 85], [293, 180]
    ])
    q = np.array([
        [211, 42], [155, 125], [100, 235],
        [235, 80], [235, 140], [295, 85], [295, 180]
    ])
    
    image = np.array(Image.open("images/toy.jpg"))
    
    height, width, _ = image.shape
    gridX = np.arange(width, dtype=np.int16)
    gridY = np.arange(height, dtype=np.int16)
    vx, vy = np.meshgrid(gridX, gridY)
    
    affine = mls_affine_deformation(vy, vx, p, q, alpha=1)
    UX =  affine[0,:,:].astype(np.float32)
    UY =  affine[1,:,:].astype(np.float32)
    aug1 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR).astype(np.uint8)
 

    similar = mls_similarity_deformation(vy, vx, p, q, alpha=1)
    UX =  similar[0,:,:].astype(np.float32)
    UY =  similar[1,:,:].astype(np.float32)
    aug2 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR).astype(np.uint8)
 

    rigid = mls_rigid_deformation(vy, vx, p, q, alpha=1)
    UX =  rigid[0,:,:].astype(np.float32)
    UY =  rigid[1,:,:].astype(np.float32)
    aug3 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR).astype(np.uint8)
 

    fig, ax = plt.subplots(1, 4, figsize=(12, 4))
    ax[0].imshow(image)
    ax[0].set_title("Original Image")    
    ax[1].imshow(aug1)
    ax[1].set_title("Affine Deformation")
    ax[2].imshow(aug2)
    ax[2].set_title("Similarity Deformation")
    ax[3].imshow(aug3)
    ax[3].set_title("Rigid Deformation")

    for x in ax.flat:
        x.axis("off")

    plt.tight_layout(w_pad=0.1)
    plt.show()


def demo_torch():


    image = cv2.imread("all.png")
    input_points = np.load("input_points.npy")
    output_points = np.load("output_points.npy")

    p = torch.from_numpy(input_points).to(device)
    q = torch.from_numpy(output_points).to(device)
    """
    p = torch.from_numpy(np.array([
        [155, 30], [155, 125], [155, 225],
        [235, 100], [235, 160], [295, 85], [293, 180]
    ])).to(device)
    q = torch.from_numpy(np.array([
        [211, 42], [155, 125], [100, 235],
        [235, 80], [235, 140], [295, 85], [295, 180]
    ])).to(device)
    """
    
    #image = cv2.imread("images/toy.jpg")
    image = torch.from_numpy(image).to(device)
    #image = torch.from_numpy(np.array(Image.open("images/toy.jpg"))).to(device)
    
    height, width, _ = image.shape
    gridX = torch.arange(width, dtype=torch.int16).to(device)
    gridY = torch.arange(height, dtype=torch.int16).to(device)
    vy, vx = torch.meshgrid(gridX, gridY)
    # !!! Pay attention !!!: the shape of returned tensors are different between numpy.meshgrid and torch.meshgrid
    vy, vx = vy.transpose(0, 1), vx.transpose(0, 1)
    
    affine = mls_affine_deformation_pt(vy, vx, p, q, alpha=1)
    aug1 = torch.ones_like(image).to(device)
    aug1[vx.long(), vy.long()] = image[tuple(affine)]
    aug1 = aug1.cpu().numpy()
    cv2.imwrite("aug1.jpg", aug1)

    similar = mls_similarity_deformation_pt(vy, vx, p, q, alpha=1)
    aug2 = torch.ones_like(image).to(device)
    aug2[vx.long(), vy.long()] = image[tuple(similar)]
    aug2 = aug2.cpu().numpy()
    cv2.imwrite("aug2.jpg", aug2)

    rigid = mls_rigid_deformation_pt(vy, vx, p, q, alpha=1)
    aug3 = torch.ones_like(image).to(device)
    aug3[vx.long(), vy.long()] = image[tuple(rigid)]
    aug3 = aug3.cpu().numpy()
    cv2.imwrite("aug3.jpg", aug3)

    """
    fig, ax = plt.subplots(1, 4, figsize=(12, 4))
    ax[0].imshow(image)
    ax[0].set_title("Original Image")    
    ax[1].imshow(aug1.cpu().numpy())
    ax[1].set_title("Affine Deformation")
    ax[2].imshow(aug2.cpu().numpy())
    ax[2].set_title("Similarity Deformation")
    ax[3].imshow(aug3.cpu().numpy())
    ax[3].set_title("Rigid Deformation")

    for x in ax.flat:
        x.axis("off")

    plt.tight_layout(w_pad=0.1)
    plt.show()
    """


def demo2():
    """ Smiled Monalisa """
    np.random.seed(1234)
    
    image = np.array(Image.open("images/monalisa.jpg"))
    height, width, _ = image.shape
    
    # Define deformation grid
    gridX = np.arange(width, dtype=np.int16)
    gridY = np.arange(height, dtype=np.int16)
    vx, vy = np.meshgrid(gridX, gridY)

    # ================ Control points group 1 (manually specified) ==================
    """
    p1 = np.array([[0, 0], [517, 0], [0, 798], [517, 798],
        [140, 186], [135, 295], [181, 208], [181, 261], [203, 184], [202, 304], [225, 213], 
        [225, 243], [244, 211], [244, 253], [254, 195], [281, 232], [252, 285]
    ])
    q1 = np.array([[0, 0], [517, 0], [0, 798], [517, 798],
        [140, 186], [135, 295], [181, 208], [181, 261], [203, 184], [202, 304], [225, 213], 
        [225, 243], [238, 207], [237, 261], [253, 199], [281, 232], [249, 279]
    ])
    """
    p1 = np.load("input_points_m.npy")
    q1 = np.load("output_points_m.npy")

    rigid1 = mls_rigid_deformation(vy, vx, p1, q1, alpha=1)
    UX =  rigid1[0,:,:].astype(np.float32)
    UY =  rigid1[1,:,:].astype(np.float32)
    aug1 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
 
    """
    # ====================== Control points group 1 (random) =======================
    p2 = np.stack((
        np.random.randint(0, height, size=13), 
        np.random.randint(0, width, size=13),
    ), axis=1)
    q2 = p2 + np.random.randint(-20, 20, size=p2.shape)

    rigid2 = mls_rigid_deformation(vy, vx, p2, q2, alpha=1)
    UX =  rigid2[0,:,:].astype(np.float32)
    UY =  rigid2[1,:,:].astype(np.float32)
    aug2 = cv2.remap(image, UX, UY, interpolation=cv2.INTER_LINEAR)
    """


    fig, ax = plt.subplots(1, 2, figsize=(13, 6))
    ax[0].imshow(image)
    ax[0].set_title("Original Image")
    ax[1].imshow(aug1)
    ax[1].set_title("Manually specified control points")
    #ax[2].imshow(aug2)
    #ax[2].set_title("Random control points")

    for x in ax.flat:
        x.axis("off")
    
    plt.tight_layout(w_pad=1.0, h_pad=1.0)
    plt.show()


def read_tif(frame):
    image_pil = Image.open("images/train-volume.tif")
    image_pil.seek(frame)
    image = np.array(image_pil)
    label_pil = Image.open("images/train-labels.tif")
    label_pil.seek(frame)
    label = np.array(label_pil)

    return image, label


def demo3():
    image, label = read_tif(1)
    image = np.pad(image, ((30, 30), (30, 30)), mode='symmetric')
    label = np.pad(label, ((30, 30), (30, 30)), mode='symmetric')

    height, width = image.shape
    gridX = np.arange(width, dtype=np.int16)
    gridY = np.arange(height, dtype=np.int16)
    vy, vx = np.meshgrid(gridX, gridY)

    def augment(p, q, mode='affine'):
        if mode.lower() == 'affine':
            transform = mls_affine_deformation(vy, vx, p, q, alpha=1)
        elif mode.lower() == 'similar':
            transform = mls_similarity_deformation(vy, vx, p, q, alpha=1)
        elif mode.lower() == 'rigid':
            transform = mls_rigid_deformation(vy, vx, p, q, alpha=1)
        else:
            raise ValueError

        aug_img = np.ones_like(image)
        aug_img[vx, vy] = image[tuple(transform)]
        aug_lab = np.ones_like(label)
        aug_lab[vx, vy] = label[tuple(transform)]

        return aug_img, aug_lab

    fig, ax = plt.subplots(2, 4, figsize=(12, 6))
    ax[0, 0].imshow(image, cmap='gray')
    ax[0, 0].set_title("Original Image")
    ax[1, 0].imshow(label, cmap='gray')
    ax[1, 0].set_title("Original Label")
    
    np.random.seed(1234)
    p = np.c_[np.random.randint(0, height, size=32), np.random.randint(0, width, size=32)]
    q = p + np.random.randint(-15, 15, size=p.shape)
    q[:, 0] = np.clip(q[:, 0], 0, height)
    q[:, 1] = np.clip(q[:, 1], 0, width)
    p = np.r_[p, np.array([[0, 0], [0, width - 1], [height - 1, 0], [height - 1, width - 1]])]  # fix corner points
    q = np.r_[q, np.array([[0, 0], [0, width - 1], [height - 1, 0], [height - 1, width - 1]])]  # fix corner points

    for i, mode in enumerate(['Affine', 'Similar', 'Rigid']):
        aug_img, aug_lab = augment(p, q, mode)
        ax[0, i + 1].imshow(aug_img, cmap='gray')
        ax[0, i + 1].set_title(f"{mode} Deformated Image")
        ax[1, i + 1].imshow(aug_lab, cmap='gray')
        ax[1, i + 1].set_title(f"{mode} Deformated Label")

    for x in ax.flat:
        x.axis('off')

    plt.tight_layout(w_pad=1.0, h_pad=1.0)
    plt.show()


def benchmark_numpy(image, p, q):
    height, width = image.shape[:2]

    # Define deformation grid
    gridX = np.arange(width, dtype=np.int16)
    gridY = np.arange(height, dtype=np.int16)
    vy, vx = np.meshgrid(gridX, gridY)

    rigid = mls_rigid_deformation(vy, vx, p, q, alpha=1)
    aug = np.ones_like(image)
    aug[vx, vy] = image[tuple(rigid)]
    return aug


def benchmark_torch(image, p, q):
    height, width = image.shape[:2]
    device = image.device

    # Define deformation grid
    gridX = torch.arange(width, dtype=torch.int16).to(device)
    gridY = torch.arange(height, dtype=torch.int16).to(device)
    vy, vx = torch.meshgrid(gridX, gridY)

    rigid = mls_rigid_deformation_pt(vy, vx, p, q, alpha=1)
    aug = torch.ones_like(image).to(device)
    aug[vx.long(), vy.long()] = image[rigid[0], rigid[1]]
    return aug


def run_benckmark(i):
    sizes = [   # (height, width)
        (100, 100),
        (500, 500),
        (500, 500),
        (500, 500),
        (1000, 1000),
        (2000, 2000),
    ]
    num_pts = [16, 16, 32, 64, 64, 64]

    times = []
    for _ in range(3):
        image = np.random.randint(0, 256, sizes[i])
        height, width = image.shape[:2]
        p = np.stack((
            np.random.randint(0, height, size=num_pts[i]), 
            np.random.randint(0, width, size=num_pts[i]),
        ), axis=1)
        q = p + np.random.randint(-20, 20, size=p.shape)

        start = time.time()
        _ = benchmark_numpy(image, p, q)
        elapse = time.time() - start
        times.append(elapse)
    print("Time (numpy):", sum(times) / len(times))

    times = []
    for _ in range(3):
        image = torch.randint(0, 256, sizes[i]).to(device)
        height, width = image.shape[:2]
        p = torch.stack((
            torch.randint(0, height, size=(num_pts[i],)),
            torch.randint(0, width, size=(num_pts[i],)),
        ), dim=1).to(device)
        q = p + torch.randint(-20, 20, size=p.shape).to(device)

        start = time.time()
        _ = benchmark_torch(image, p, q)
        elapse = time.time() - start
        times.append(elapse)
    print("Time (torch):", sum(times) / len(times))



def test():
    srcImg = cv2.imread("all.png")
    #srcImg = srcImg[:720,:720,:]
    input_points = np.load("input_points_271.npy")
    output_points = np.load("output_points_271.npy")
    
    src = srcImg.copy()
    for i in range(0,input_points.shape[0],1):
        cv2.circle(src,(int(input_points[i,0]), int(input_points[i,1])), 1, (0,0,255), cv2.MARKER_SQUARE)
        cv2.circle(src,(int(output_points[i,0]), int(output_points[i,1])), 1, (0,255,255), cv2.MARKER_CROSS)
    cv2.imwrite("src.jpg", src)
    
    aug = MLS(srcImg, input_points, output_points)
    #aug = MLS_TORCH(srcImg, input_points, output_points)

    
    for i in range(0,input_points.shape[0],1):
        cv2.circle(aug,(int(input_points[i,0]), int(input_points[i,1])), 1, (0,0,255), cv2.MARKER_SQUARE)
        cv2.circle(aug,(int(output_points[i,0]), int(output_points[i,1])), 1, (0,255,255), cv2.MARKER_CROSS)
    cv2.imwrite("aug.jpg", aug)
    
 




if __name__ == "__main__":
    #demo()
    #demo_torch()
    demo2()
    # demo3()

    # run_benckmark(i=0)
    #test()
