import gradio as gr
import bpy
from tqdm import tqdm
from math import pi
import tempfile
import molecularnodes as mn
import os

def enable_GPUS():
    bpy.data.scenes[0].render.engine = "CYCLES" #"CYCLES"
    # Set the device_type
    bpy.context.preferences.addons[
        "cycles"
    ].preferences.compute_device_type = "CUDA"  # or "OPENCL"

    # Set the device and feature set
    bpy.context.scene.cycles.device = "GPU"

    for scene in bpy.data.scenes:
        scene.cycles.device = "GPU"

    bpy.context.preferences.addons["cycles"].preferences.get_devices()
    print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type)
    for d in bpy.context.preferences.addons["cycles"].preferences.devices:
        d["use"] = True  # Using all devices, include GPU and CPU
        print(d["name"])


enable_GPUS()

window = bpy.context.window
screen = window.screen

style = 'cartoon'

nodes_to_append = ["MN_color_set", 
                   "MN_color_common", 
                   "MN_color_attribute_random",
                   mn.nodes.styles_mapping[style]]

def get_areas(type):
    return [area for area in screen.areas if area.type == type]

def get_regions(areas):
    return [region for region in areas[0].regions if region.type == 'WINDOW']

for node in nodes_to_append:
    bpy.ops.wm.append(
                'INVOKE_DEFAULT',
                directory = os.path.join(mn.nodes.mn_data_file, 'NodeTree'), 
                filename = node, 
                link = False
            )

def generate(progress=gr.Progress(track_tqdm=True)):
    area_type = 'VIEW_3D' 
    areas  = get_areas(area_type)
    with bpy.context.temp_override(window=window, area=areas[0], region=get_regions(areas)[0], screen=screen):
        
        
        for obj in bpy.context.scene.objects:
            if obj.type == 'MESH':
                bpy.data.objects.remove(obj, do_unlink=True)

        molecule = mn.load.molecule_rcsb("7TYG", starting_style=style, center_molecule=True)
        molecule.select_set(True)
        bpy.context.view_layer.objects.active = molecule

        bpy.ops.view3d.camera_to_view_selected()
        camera = bpy.data.objects["Camera"]
        camera.data.dof.use_dof = True
        camera.data.dof.focus_distance = 5
        camera.data.dof.aperture_fstop = 4
        camera.data.angle = pi / 3
        camera.data.type = "PERSP"

        with tempfile.NamedTemporaryFile(suffix=".JPEG", delete=False) as f:
            bpy.context.scene.render.resolution_y = 400
            bpy.context.scene.render.resolution_x = 600
            bpy.context.scene.render.image_settings.file_format = "JPEG"
            bpy.context.scene.render.filepath = f.name

            with tqdm() as pbar:

                def elapsed(dummy):
                    pbar.update()

                bpy.app.handlers.render_stats.append(elapsed)
                bpy.context.scene.frame_set(1)
                bpy.context.scene.frame_current = 1
                bpy.ops.render.render(animation=False, write_still=True)

                bpy.data.images["Render Result"].save_render(
                    filepath=bpy.context.scene.render.filepath
                )
                bpy.app.handlers.render_stats.clear()
                return f.name

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            render_btn = gr.Button("Render")
        with gr.Column(scale=3):
            image = gr.Image(type="filepath")

    render_btn.click(
        generate,
        outputs=[image],
    )

demo.queue(concurrency_count=1)
demo.launch(debug=True, inline=True)