import sys
from subprocess import check_call
import tempfile

from os.path import basename, splitext, join
from io import BytesIO

import numpy as np
from scipy.spatial import KDTree
from PIL import Image

import torch
import torch.nn.functional as F
from torchvision.transforms.functional import to_tensor, to_pil_image
from einops import rearrange
import gradio as gr
from huggingface_hub import hf_hub_download

from extern.ZoeDepth.zoedepth.utils.misc import colorize

from gradio_model3dgscamera import Model3DGSCamera

def download_models():
    models = [
        {
            'repo': 'stabilityai/sd-vae-ft-mse',
            'sub': None,
            'dst': 'checkpoints/sd-vae-ft-mse',
            'files': ['config.json', 'diffusion_pytorch_model.safetensors'],
            'token': None
        },
        {
            'repo': 'lambdalabs/sd-image-variations-diffusers',
            'sub': 'image_encoder',
            'dst': 'checkpoints',
            'files': ['config.json', 'pytorch_model.bin'],
            'token': None
        },
        {
            'repo': 'Sony/genwarp',
            'sub': 'multi1',
            'dst': 'checkpoints',
            'files': ['config.json', 'denoising_unet.pth', 'pose_guider.pth', 'reference_unet.pth'],
            'token': None
        }
    ]

    for model in models:
        for file in model['files']:
            hf_hub_download(
                repo_id=model['repo'],
                subfolder=model['sub'],
                filename=file,
                local_dir=model['dst'],
                token=model['token']
            )

# Setup.
download_models()

mde = torch.hub.load(
    './extern/ZoeDepth',
    'ZoeD_N',
    source='local',
    pretrained=True,
    trust_repo=True
)

import spaces

check_call([
    sys.executable, '-m', 'pip', 'install',
    'extern/splatting-0.0.1-py3-none-any.whl'
])

from genwarp import GenWarp
from genwarp.ops import (
    camera_lookat, get_projection_matrix, get_viewport_matrix
)

# GenWarp
genwarp_cfg = dict(
    pretrained_model_path='checkpoints',
    checkpoint_name='multi1',
    half_precision_weights=True
)
genwarp_nvs = GenWarp(cfg=genwarp_cfg, device='cpu')

# Fixed parameters.
IMAGE_SIZE = 512
NEAR, FAR = 0.01, 100
FOVY = np.deg2rad(55)
PROJ_MTX = get_projection_matrix(
    fovy=torch.ones(1) * FOVY,
    aspect_wh=1.,
    near=NEAR,
    far=FAR
)
VIEW_MTX = camera_lookat(
    torch.tensor([[0., 0., 0.]]),
    torch.tensor([[0., 0., 1.]]),
    torch.tensor([[0., -1., 0.]])
)
VIEWPORT_MTX = get_viewport_matrix(
    IMAGE_SIZE, IMAGE_SIZE,
    batch_size=1
)

# Crop the image to the shorter side.
def crop(img: Image) -> Image:
    W, H = img.size
    if W < H:
        left, right = 0, W
        top, bottom = np.ceil((H - W) / 2.), np.floor((H - W) / 2.) + W
    else:
        left, right = np.ceil((W - H) / 2.), np.floor((W - H) / 2.) + H
        top, bottom = 0, H
    img = img.crop((left, top, right, bottom))
    img = img.resize((IMAGE_SIZE, IMAGE_SIZE))
    return img

def save_as_splat(
    filepath: str,
    xyz: np.ndarray,
    rgb: np.ndarray
):
    # To gaussian splat
    inv_sigmoid = lambda x: np.log(x / (1 - x))
    dist2 = np.clip(calc_dist2(xyz), a_min=0.0000001, a_max=None)
    scales = np.repeat(np.log(np.sqrt(dist2))[..., np.newaxis], 3, axis=1)
    rots = np.zeros((xyz.shape[0], 4))
    rots[:, 0] = 1
    opacities = inv_sigmoid(0.1 * np.ones((xyz.shape[0], 1)))

    sorted_indices = np.argsort((
        -np.exp(np.sum(scales, axis=-1, keepdims=True))
        / (1 + np.exp(-opacities))
    ).squeeze())

    buffer = BytesIO()
    for idx in sorted_indices:
        position = xyz[idx]
        scale = np.exp(scales[idx]).astype(np.float32)
        rot = rots[idx].astype(np.float32)
        color = np.concatenate(
            (rgb[idx], 1 / (1 + np.exp(-opacities[idx]))),
            axis=-1
        )
        buffer.write(position.tobytes())
        buffer.write(scale.tobytes())
        buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
        buffer.write(
            ((rot / np.linalg.norm(rot)) * 128 + 128)
            .clip(0, 255)
            .astype(np.uint8)
            .tobytes()
        )

    with open(filepath, "wb") as f:
        f.write(buffer.getvalue())

def calc_dist2(points: np.ndarray):
    dists, _ = KDTree(points).query(points, k=4)
    mean_dists = (dists[:, 1:] ** 2).mean(1)
    return mean_dists

def unproject(depth):
    H, W = depth.shape[2:4]
    mean_depth = depth.mean(dim=(2, 3)).squeeze().item()

    # Matrices.
    viewport_mtx = VIEWPORT_MTX.to(depth)
    proj_mtx = PROJ_MTX.to(depth)
    view_mtx = VIEW_MTX.to(depth)
    scr_mtx = (viewport_mtx @ proj_mtx).to(depth)

    grid = torch.stack(torch.meshgrid(
        torch.arange(W), torch.arange(H), indexing='xy'), dim=-1
    ).to(depth)[None]  # BHW2

    screen = F.pad(grid, (0, 1), 'constant', 0)
    screen = F.pad(screen, (0, 1), 'constant', 1)
    screen_flat = rearrange(screen, 'b h w c -> b (h w) c')

    eye = screen_flat @ torch.linalg.inv_ex(
        scr_mtx.float()
    )[0].mT.to(depth)
    eye = eye * rearrange(depth, 'b c h w -> b (h w) c')
    eye[..., 3] = 1

    points = eye @ torch.linalg.inv_ex(view_mtx.float())[0].mT.to(depth)
    points = points[0, :, :3]

    # Translate to the origin.
    points[..., 2] -= mean_depth
    camera_pos = (0, 0, -mean_depth)

    return points, camera_pos

def view_from_rt(position, rotation):
    t = np.array(position)
    euler = np.array(rotation)

    cx = np.cos(euler[0])
    sx = np.sin(euler[0])
    cy = np.cos(euler[1])
    sy = np.sin(euler[1])
    cz = np.cos(euler[2])
    sz = np.sin(euler[2])
    R = np.array([
        cy * cz + sy * sx * sz,
        -cy * sz + sy * sx * cz,
        sy * cx,
        cx * sz,
        cx * cz,
        -sx,
        -sy * cz + cy * sx * sz,
        sy * sz + cy * sx * cz,
        cy * cx
    ])
    view_mtx = np.array([
        [R[0], R[1], R[2], 0],
        [R[3], R[4], R[5], 0],
        [R[6], R[7], R[8], 0],
        [
            -t[0] * R[0] - t[1] * R[3] - t[2] * R[6],
            -t[0] * R[1] - t[1] * R[4] - t[2] * R[7],
            -t[0] * R[2] - t[1] * R[5] - t[2] * R[8],
            1
        ]
    ]).T

    B = np.array([
        [1, 0, 0, 0],
        [0, -1, 0, 0],
        [0, 0, -1, 0],
        [0, 0, 0, 1]
    ])
    return B @ view_mtx


with tempfile.TemporaryDirectory() as tmpdir:
    with gr.Blocks(
        title='GenWarp Demo',
        css='img {display: inline;}'
    ) as demo:
        # Internal states.
        image = gr.State()
        depth = gr.State()

        # Callbacks
        @spaces.GPU()
        def cb_mde(image_file: str):
            # Load an image.
            image_pil = crop(Image.open(image_file).convert('RGB'))
            image = to_tensor(image_pil)[None].detach()
            # Get depth.
            depth = mde.cuda().infer(image.cuda()).cpu().detach()
            depth_pil = to_pil_image(colorize(depth[0]))
            return image_pil, depth_pil, image, depth

        @spaces.GPU()
        def cb_3d(image_file, image, depth):
            # Unproject.
            xyz, camera_pos = unproject(depth.cuda())
            xyz = xyz.cpu().detach().numpy()
            # Save as a splat.
            ## Output filename.
            splat_file = join(
                tmpdir, f'./{splitext(basename(image_file))[0]}.splat')
            rgb = rearrange(image, 'b c h w -> b (h w) c')[0].numpy()
            save_as_splat(splat_file, xyz, rgb)
            return splat_file, camera_pos, (0, 0, 0)

        @spaces.GPU()
        def cb_generate(viewer, image, depth):
            if depth is None:
                gr.Error('Image and Depth are not set. Try again.')
                return None, None

            mean_depth = depth.mean(dim=(2, 3)).squeeze().item()
            src_view_mtx = camera_lookat(
                torch.tensor([[0., 0., -mean_depth]]),
                torch.tensor([[0., 0., 0.]]),
                torch.tensor([[0., -1., 0.]])
            ).to(depth)
            tar_camera_pos, tar_camera_rot = viewer[1:3]
            tar_view_mtx = torch.from_numpy(view_from_rt(
                tar_camera_pos, tar_camera_rot
            ))
            rel_view_mtx = (
                tar_view_mtx @ torch.linalg.inv(src_view_mtx.double())
            ).half().cuda()
            proj_mtx = PROJ_MTX.half().cuda()

            # GenWarp.
            renders = genwarp_nvs.to('cuda')(
                src_image=image.half().cuda(),
                src_depth=depth.half().cuda(),
                rel_view_mtx=rel_view_mtx,
                src_proj_mtx=proj_mtx,
                tar_proj_mtx=proj_mtx
            )
            warped_pil = to_pil_image(renders['warped'].cpu()[0])
            synthesized_pil = to_pil_image(renders['synthesized'].cpu()[0])

            return warped_pil, synthesized_pil

        def process_example(image_file):
            gr.Error('')
            image_pil, depth_pil, image, depth = cb_mde(image_file)
            viewer = cb_3d(image_file, image, depth)
            # Fixed angle for examples.
            viewer = (viewer[0], (-2.020, -0.727, -5.236), (-0.132, 0.378, 0.0))
            warped_pil, synthsized_pil = cb_generate(
                viewer, image, depth
            )
            return (
                image_pil, depth_pil, viewer,
                warped_pil, synthsized_pil,
                None, None  # Clear internal states.
            )

        # Blocks.
        gr.Markdown(
            """
            # GenWarp: Single Image to Novel Views with Semantic-Preserving Generative Warping
            [![Project Site](https://img.shields.io/badge/Project-Web-green)](https://genwarp-nvs.github.io/) &nbsp;
            [![Spaces](https://img.shields.io/badge/Spaces-Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/Sony/GenWarp) &nbsp;
            [![Github](https://img.shields.io/badge/Github-Repo-orange?logo=github)](https://github.com/sony/genwarp/) &nbsp;
            [![Models](https://img.shields.io/badge/Models-checkpoints-blue?logo=huggingface)](https://huggingface.co/Sony/genwarp) &nbsp;
            [![arXiv](https://img.shields.io/badge/arXiv-2405.17251-red?logo=arxiv)](https://arxiv.org/abs/2405.17251)

            ## Introduction
            This is an official demo for the paper "[GenWarp: Single Image to Novel Views with Semantic-Preserving Generative Warping](https://genwarp-nvs.github.io/)". Genwarp can generate novel view images from a single input conditioned on camera poses. In this demo, we offer a basic use of inference of the model. For detailed information, please refer to the [paper](https://arxiv.org/abs/2405.17251).

            ## How to Use

            ### Try examples
            - Examples are in the bottom section of the page

            ### Upload your own images
            1. Upload a reference image to "Reference Input"
            2. Move the camera to your desired view in "Unprojected 3DGS" 3D viewer
            3. Hit "Generate a novel view" button and check the result

            ## Tips
            - This model is mainly trained for indoor/outdoor scenery. It might not work well for object-centric inputs. For details on training the model, please check our [paper](https://arxiv.org/abs/2405.17251).
            - Extremely large camera movement from the input view might cause low performance results due to the unexpected deviation from the training distribution, which is not the scope of this model. Instead, you can feed the generation result for the small camera movement repeatedly and progressively move towards a desired view.
            - 3D viewer might take some time to update especially when trying different images back to back. Wait until it fully updates to the new image.

            """
        )
        file = gr.File(label='Reference Input', file_types=['image'])
        with gr.Row():
            image_widget = gr.Image(
                label='Reference View', type='filepath',
                interactive=False
            )
            depth_widget = gr.Image(label='Estimated Depth', type='pil')
            viewer = Model3DGSCamera(
                label = 'Unprojected 3DGS',
                width=IMAGE_SIZE,
                height=IMAGE_SIZE,
                camera_width=IMAGE_SIZE,
                camera_height=IMAGE_SIZE,
                camera_fx=IMAGE_SIZE / (np.tan(FOVY / 2.)) / 2.,
                camera_fy=IMAGE_SIZE / (np.tan(FOVY / 2.)) / 2.,
                camera_near=NEAR,
                camera_far=FAR
            )
        button = gr.Button('Generate a novel view', size='lg', variant='primary')
        with gr.Row():
            warped_widget = gr.Image(
                label='Warped Image', type='pil', interactive=False
            )
            gen_widget = gr.Image(
                label='Generated View', type='pil', interactive=False
            )
        examples = gr.Examples(
            examples=[
                './assets/pexels-heyho-5998120_19mm.jpg',
                './assets/pexels-itsterrymag-12639296_24mm.jpg'
            ],
            fn=process_example,
            inputs=file,
            outputs=[image_widget, depth_widget, viewer,
                     warped_widget, gen_widget,
                     image, depth]
        )

        # Events
        file.upload(
            fn=cb_mde,
            inputs=file,
            outputs=[image_widget, depth_widget, image, depth]
        ).success(
            fn=cb_3d,
            inputs=[image_widget, image, depth],
            outputs=viewer
        )
        button.click(
            fn=cb_generate,
            inputs=[viewer, image, depth],
            outputs=[warped_widget, gen_widget]
        )
        # To re-calculate the uncached depth for examples in background.
        examples.load_input_event.success(
            fn=lambda x: cb_mde(x)[2:4],
            inputs=file,
            outputs=[image, depth]
        )

    if __name__ == '__main__':
        demo.launch()