import gradio as gr
import pickle
# from datasets import load_from_disk
from plaid.containers.sample import Sample
# import pyvista as pv

import pyrender
import trimesh
import matplotlib.pyplot as plt

import os
# switch to "osmesa" or "egl" before loading pyrender
os.environ["PYOPENGL_PLATFORM"] = "egl"

import numpy as np

# FOLDER = "plot"

# dataset = load_from_disk("Rotor37")

field_names_train = ["Temperature", "Pressure", "Density"]#pickle.loads(dataset[0]["sample"]).get_field_names()
field_names_test  = []







def sample_info(sample_id_str, fieldn):

    # sample_id = int(sample_id_str)
    # plaid_sample = Sample.load_from_dir(f"Rotor37/dataset/samples/sample_"+str(sample_id_str).zfill(9))

    str__ = f"loading sample {sample_id_str}"
    
    
    # generate mesh
    sphere = trimesh.creation.icosphere(subdivisions=4, radius=0.8)
    sphere.vertices+=1e-2*np.random.randn(*sphere.vertices.shape)
    mesh = pyrender.Mesh.from_trimesh(sphere, smooth=False)
    
    # compose scene
    scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[0, 0, 0])
    camera = pyrender.PerspectiveCamera( yfov=np.pi / 3.0)
    light = pyrender.DirectionalLight(color=[1,1,1], intensity=2e3)
    
    scene.add(mesh, pose=  np.eye(4))
    scene.add(light, pose=  np.eye(4))
    
    c = 2**-0.5
    scene.add(camera, pose=[[ 1,  0,  0,  0],
                            [ 0,  c, -c, -2],
                            [ 0,  c,  c,  2],
                            [ 0,  0,  0,  1]])
    
    # render scene
    r = pyrender.OffscreenRenderer(512, 512)
    color, _ = r.render(scene)
    # color = np.random.rand(512, 512)
    
    plt.figure(figsize=(8,8))
    plt.imshow(color)
    
    plt.savefig("test.png")



    return str__, "test.png"
    # return str__, str__


if __name__ == "__main__":

    with gr.Blocks() as demo:
        d1 = gr.Slider(0, 999, value=0, label="Training sample id", info="Choose between 0 and 999")
        d2 = gr.Dropdown(field_names_train, value=field_names_train[0], label="Field name")


        output1 = gr.Text(label="Training sample info")
        # output2 = gr.Text(label="Training sample visualization")
        output2 = gr.Image(label="Training sample visualization")

        # d1.input(update_second, d1, d2)

        d1.input(sample_info, [d1, d2], [output1, output2])
        d2.input(sample_info, [d1, d2], [output1, output2]) 


    demo.launch()