import torch
import cheetah
import gradio as gr
import os

# Define the accelerator segment
segment = cheetah.Segment(
    elements=[
        cheetah.Drift(length=torch.tensor(0.175)),
        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM1"),
        cheetah.Drift(length=torch.tensor(0.428)),
        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM2"),
        cheetah.Drift(length=torch.tensor(0.204)),
        cheetah.VerticalCorrector(length=torch.tensor(0.02), name="AREAMCVM1"),
        cheetah.Drift(length=torch.tensor(0.204)),
        cheetah.Quadrupole(length=torch.tensor(0.122), name="AREAMQZM3"),
        cheetah.Drift(length=torch.tensor(0.179)),
        cheetah.HorizontalCorrector(length=torch.tensor(0.02), name="AREAMCHM1"),
        cheetah.Drift(length=torch.tensor(0.45)),
        cheetah.Screen(name="AREABSCR1"),
    ]
)

# Generate the mesh
mesh, _ = segment.to_mesh(
    cuteness={cheetah.HorizontalCorrector: 2.0, cheetah.VerticalCorrector: 2.0}
)

# Export the mesh to a .glb file
output_file = "ares_mesh.glb"
mesh.export(file_obj=output_file, file_type="glb")

# Define the Gradio interface
def display_3d_model():
    return output_file

with gr.Blocks() as demo:
    gr.Markdown("# 3D Visualization of Cheetah Accelerator Segment")
    gr.Markdown("This app renders a 3D model of the accelerator segment defined using the Cheetah library.")
    gr.File(file_types=[".glb"], label="3D Model", value=output_file)
    gr.Model3D(value=output_file, label="3D Visualization")

# Launch the Gradio app
demo.launch()