import os
import gradio as gr
from PIL import Image
import torch
import matplotlib.pyplot as plt
import imageio
import numpy as np
import math
import argparse
import tempfile

import torch
import base64
import io
import os
from typing import Union


from shap_e.diffusion.sample import sample_latents
from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
from shap_e.models.download import load_model, load_config
from shap_e.util.notebooks import create_pan_cameras, decode_latent_images, decode_latent_mesh

from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
from shap_e.models.transmitter.base import Transmitter, VectorDecoder
from shap_e.util.collections import AttrDict

import trimesh


state = ""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

css = '''
    .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
    .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
    #component-4, #component-3, #component-10{min-height: 0}
'''

def set_state(s):
    print(s)
    global state
    state = s

def get_state():
    return state

def to_video(frames: list[Image.Image], fps: int = 5) -> str:
        out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
        writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
        for frame in frames:
            writer.append_data(np.asarray(frame))
        writer.close()
        return out_file.name

def generate_3D(input, grid_size=64):
    set_state('Entered generate function...')

    # if input is a string, it's a text prompt
    xm = load_model('transmitter', device=device)
    diffusion = diffusion_from_config(load_config('diffusion'))
    batch_size = 4

    if isinstance(input, np.ndarray):
        input = Image.fromarray(input)

    if isinstance(input, Image.Image):
        input = prepare_img(input)
        model = load_model('image300M', device=device)
        guidance_scale = 3.0
        model_kwargs = dict(images=[input] * batch_size)
    else:
        model = load_model('text300M', device=device)
        guidance_scale = 15.0
        model_kwargs = dict(texts=[input] * batch_size)

    print(input)

    latents = sample_latents(
        batch_size=batch_size,
        model=model,
        diffusion=diffusion,
        guidance_scale=guidance_scale,
        model_kwargs=model_kwargs,
        progress=True,
        clip_denoised=True,
        use_fp16=True,
        use_karras=True,
        karras_steps=64,
        sigma_min=1e-3,
        sigma_max=160,
        s_churn=0,
    )

    render_mode = 'stf' # you can change this to 'stf'
    size = grid_size # this is the size of the renders; higher values take longer to render.

    cameras = create_pan_cameras(size, device)

    with open(f'/tmp/mesh.ply', 'wb') as f:
        decode_latent_mesh(xm, latents[0]).tri_mesh().write_ply(f)


    set_state('Converting to point cloud...')
    # pc = sampler.output_to_point_clouds(samples)[0]

    set_state('Converting to mesh...')
    # save_ply(pc, 'output/mesh.ply', grid_size)

    set_state('')

    images = decode_latent_images(xm, latents[0], cameras, rendering_mode=render_mode)


    return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), to_video(images), gr.update(value=['/tmp/mesh.glb', '/tmp/mesh.ply'], visible=True)

def prepare_img(img):

    w, h = img.size
    if w > h:
        img = img.crop((w - h) / 2, 0, w - (w - h) / 2, h)
    else:
        img = img.crop((0, (h - w) / 2, w, h - (h - w) / 2))

    # resize to 256x256
    img = img.resize((256, 256))

    return img


def ply_to_glb(ply_file, glb_file):
    mesh = trimesh.load(ply_file)

    # Save the mesh as a glb file using Trimesh
    mesh.export(glb_file, file_type='glb')

    return glb_file


block = gr.Blocks().queue(max_size=250, concurrency_count=6)
with block:
    with gr.Box():
        if(not torch.cuda.is_available()):
            top_description = gr.HTML(f'''
                <div style="text-align: center; max-width: 650px; margin: 0 auto;">
                <div>
                    <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"
                        style="margin: auto; max-width: 7rem;">
                    <br />
                    <h1 style="font-weight: 900; font-size: 2.5rem;">
                    Shap-E Web UI
                    </h1>
                </div>
                <h3 style="font-weight: 900; font-size: 1.5rem;">
                If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>!
                </h3>
                <br />
                <p style="margin-bottom: 10px; font-size: 94%">
                Generate 3D Assets in 1 minute with a prompt or image!
                Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation
                </p>
                <br />
                <p>There's only one step left before you can train your model: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4 GPU</b> to it (via the Settings tab)</a> and run the training below. Other GPUs are not compatible for now. You will be billed by the minute from when you activate the GPU until when it is turned it off.</p>
                </div>
            ''')
        else:
            top_description = gr.HTML(f'''
                    <div style="text-align: center; max-width: 650px; margin: 0 auto;">
                    <div>
                        <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"
                            style="margin: auto; max-width: 7rem;">
                        <br />
                        <h1 style="font-weight: 900; font-size: 2.5rem;">
                        Shap-E Web UI
                        </h1>
                    </div>
                    <h3 style="font-weight: 900; font-size: 1.5rem;">
                    If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>!
                    </h3>
                    <br />
                    <p style="margin-bottom: 10px; font-size: 94%">
                    Generate 3D Assets in 1 minute with a prompt or image!
                    Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation
                    </p>
                    </div>
                ''')
    with gr.Row():
        with gr.Column():
            with gr.Tab("Text to 3D"):
                gr.Markdown("Uses Stable Diffusion to create an image from the prompt.")
                prompt = gr.Textbox(label="Prompt", placeholder="A HD photo of a Corgi")
                text_button = gr.Button(label="Generate")

            with gr.Tab("Image to 3D"):
                gr.Markdown("Best results with images of objects on an empty background.")
                input_image = gr.Image(label="Image")
                img_button = gr.Button(label="Generate")

            # with gr.Accordion("Advanced options", open=False):
            #     model = gr.Radio(["base40M", "base300M", "base1B"], label="Model", value="base1B")
            #     scale = gr.Slider(
            #         label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1
            #     )

        with gr.Column():
            model_gif = gr.Model3D(label="3D Model GIF")
            # btn_pc_to_obj = gr.Button(value="Convert to OBJ", visible=False)
            model_3d = gr.Model3D(value=None)
            file_out = gr.File(label="Files", visible=False)

    if torch.cuda.is_available():
        gr.Examples(
            examples=[
                ["a shark"],
                ["an avocado"],
            ],
            inputs=[prompt],
            outputs=[model_3d, model_gif, file_out],
            fn=generate_3D,
            cache_examples=True
        )
        gr.Examples(
            examples=[
                ["images/pumpkin.png"],
                ["images/fantasy_world.png"],
            ],
            inputs=[input_image],
            outputs=[model_3d, model_gif, file_out],
            fn=generate_3D,
            cache_examples=True
        )

    img_button.click(fn=generate_3D, inputs=[input_image], outputs=[model_3d, model_gif, file_out])
    text_button.click(fn=generate_3D, inputs=[prompt], outputs=[model_3d, model_gif, file_out])

block.launch(show_api=False)