import gradio as gr
import torch
import numpy as np
import cv2
import tempfile
from PIL import Image
import torchvision.transforms as transforms
import os
from huggingface_hub import hf_hub_download
import huggingface_hub

from models.pvsnet_model import PVSNet
import helperFunctions as helper
import parameters as params

DEVICE = params.DEVICE

def getPositionVector(x, y, z, pose_dims=3):
    if pose_dims == 3:
        vector = torch.zeros((1, 3), dtype=torch.float)
        vector[0][0] = (float(format(x, '.7f')) - (-0.1)) / (0.1 - (-0.1))
        vector[0][1] = (float(format(y, '.7f')) - (-0.1)) / (0.1 - (-0.1))
        vector[0][2] = (float(format(z, '.7f')) - (-0.1)) / (0.1 - (-0.1))
        return vector
    else:
        t_min, t_max = -0.1, 0.1
        r_min, r_max = -3, 3
        vector = torch.zeros((1, 6), dtype=torch.float)
        vector[0, 0] = (x - t_min) / (t_max - t_min)
        vector[0, 1] = (y - t_min) / (t_max - t_min)
        vector[0, 2] = (z - t_min) / (t_max - t_min)
        vector[0, 3] = (0 - r_min) / (r_max - r_min)
        vector[0, 4] = (0 - r_min) / (r_max - r_min)
        vector[0, 5] = (0 - r_min) / (r_max - r_min)
        return vector

def generateCircularTrajectory(radius, num_frames):
    angles = np.linspace(0, 2 * np.pi, num_frames, endpoint=False)
    return [[radius * np.cos(angle), radius * np.sin(angle), 0] for angle in angles]

def generateSwingTrajectory(radius, num_frames):
    angles = np.linspace(0, 2 * np.pi, num_frames, endpoint=False)
    return [[radius * np.cos(angle), 0, radius * np.sin(angle)] for angle in angles]

def create_video_from_memory(frames, fps=30):
    if not frames:
        return None
    height, width, _ = frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    out = cv2.VideoWriter(temp_video.name, fourcc, fps, (width, height))
    for frame in frames:
        out.write(frame)
    out.release()
    return temp_video.name

def process_image(img, video_type, radius, num_frames, num_loops, dataset, resolution, architecture):
    if img is None:
        return None
        
    width, height = map(int, resolution.split('x'))

    min_dim = min(img.width, img.height)
    left = (img.width - min_dim) / 2
    top = (img.height - min_dim) / 2
    right = (img.width + min_dim) / 2
    bottom = (img.height + min_dim) / 2
    img = img.crop((left, top, right, bottom))

    is_lite = (architecture == "Lite")
    pose_dims = 3 if dataset == "Blender" else 6
    
    dataset_prefix = dataset.lower()
    arch_infix = "pvsnet_lite" if is_lite else "pvsnet"
    checkpoint_name = f"checkpoint_{dataset_prefix}_{arch_infix}_{resolution}.pth"
    try:
        checkpoint_path = hf_hub_download(
            repo_id="3ZadeSSG/PVSNet", 
            filename=checkpoint_name
        )
    except huggingface_hub.utils.EntryNotFoundError:
        raise gr.Error(f"Checkpoint {checkpoint_name} not found in Hugging Face Hub! Please select a valid combination.")
    except Exception as e:
        raise gr.Error(f"Error downloading checkpoint {checkpoint_name}: {e}")

    model = PVSNet(total_image_input=params.params_number_input, pose_dims=pose_dims, height=height, width=width, is_lite=is_lite)
    
    try:
        model = helper.load_Checkpoint(checkpoint_path, model, load_cpu=True)
    except Exception as e:
        print(f"Error loading checkpoint {checkpoint_path}: {e}")
        raise gr.Error(f"Error loading checkpoint {checkpoint_path}: {e}")
        
    model.to(DEVICE)
    model.eval()

    transform = transforms.Compose([
        transforms.Resize((height, width)),
        transforms.ToTensor()
    ])
    
    img_input = img.convert('RGB')
    img_input = transform(img_input).unsqueeze(0).to(DEVICE)

    if video_type == "Circle":
        raw_traj = generateCircularTrajectory(radius, num_frames)
        trajectory = [(p[0], p[1], 0) for p in raw_traj]
    elif video_type == "Swing":
        raw_traj = generateSwingTrajectory(radius, num_frames)
        trajectory = raw_traj
    else:
         raw_traj = generateCircularTrajectory(radius, num_frames)
         trajectory = [(p[0], p[1], 0) for p in raw_traj]

    view_frames = []

    for x, y, z in trajectory:
        pos = getPositionVector(x, y, z, pose_dims=pose_dims).unsqueeze(0).to(DEVICE)
        
        with torch.no_grad():
            predicted_img = model(img_input, pos)
        
        p_img = predicted_img[0].detach().cpu().permute(1, 2, 0).numpy()
        p_img = np.clip(p_img, 0, 1)
        p_img = (p_img * 255).astype(np.uint8)
        p_img_bgr = cv2.cvtColor(p_img, cv2.COLOR_RGB2BGR)
        view_frames.append(p_img_bgr)

    view_frames = view_frames * int(num_loops)

    fps = 60
    view_video_path = create_video_from_memory(view_frames, fps=fps)

    return view_video_path

with gr.Blocks(title="PVSNet", theme="default") as demo:
    gr.Markdown(
    """
    ## PVSNet: Real-Time Position-Aware View Synthesis from Single-View Input
    * Upload an image and get a mini video showing capability of novel view synthesis.
    **Note:** Huggingface demo is running on CPU so inference speeds will be slow. Inference might take around 2-5 mins depending on resolution and model. We recomment runnning it on 256x256 and with Lite model.
    ### Head to our [Project Page](https://realistic3d-miun.github.io/PVSNet/) for more details about the models
    """)
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Input Image", height=256)
            
            with gr.Group():
                dataset_type = gr.Dropdown(["Blender", "COCO"], label="Dataset Model", value="COCO")
                resolution_type = gr.Dropdown(["256x256", "512x512"], label="Resolution", value="256x256")
                architecture_type = gr.Dropdown(["Regular", "Lite"], label="Architecture", value="Lite")
                video_type = gr.Dropdown(["Circle", "Swing"], label="Trajectory Type", value="Swing")
            
            with gr.Accordion("Advanced Settings", open=False):
                radius = gr.Slider(0.01, 0.1, value=0.06, label="Motion Radius")
                num_frames = gr.Slider(10, 120, value=60, step=1, label="Frames per Loop")
                num_loops = gr.Slider(1, 6, value=3, step=1, label="Number of Loops")

            
            submit_btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            video_output = gr.Video(label="Generated View Video", height=256)

    submit_btn.click(
        fn=process_image,
        inputs=[img_input, video_type, radius, num_frames, num_loops, dataset_type, resolution_type, architecture_type],
        outputs=[video_output]
    )

    gr.Markdown("### Example Images: Click to Load")
    
    import glob
    blender_imgs = sorted(glob.glob("./sample_images/blender/*"))
    coco_imgs = sorted(glob.glob("./sample_images/coco/*"))
    rw_imgs = sorted(glob.glob("./sample_images/real_world/*"))

    def create_grid(imgs, title, cols=4):
        gr.Markdown(title)
        image_components = []
        for i in range(0, len(imgs), cols):
            with gr.Row():
                for img_path in imgs[i:i+cols]:
                    comp = gr.Image(img_path, label=os.path.basename(img_path), height=150, interactive=False, show_label=True)
                    image_components.append((comp, img_path))
        return image_components

    with gr.Column():
        b_comps = create_grid(blender_imgs, "#### Blender Models (Loads Blender Lite 256x256 by Default)")
        c_comps = create_grid(coco_imgs, "#### COCO Models (Loads COCO 256x256 Lite by Default)")
        r_comps = create_grid(rw_imgs, "#### Real World Models (Loads COCO 256x256 Lite by Default)")

    for comp, path in b_comps:
        comp.select(fn=lambda p=path: (Image.open(p), "Blender", "256x256", "Lite"), outputs=[img_input, dataset_type, resolution_type, architecture_type])
    
    for comp, path in c_comps:
        comp.select(fn=lambda p=path: (Image.open(p), "COCO", "256x256", "Lite"), outputs=[img_input, dataset_type, resolution_type, architecture_type])

    for comp, path in r_comps:
        comp.select(fn=lambda p=path: (Image.open(p), "COCO", "256x256", "Lite"), outputs=[img_input, dataset_type, resolution_type, architecture_type])

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