import gradio as gr
import torch
import numpy as np
import plotly.graph_objects as go
from transformers import AutoModel

# 1. Загрузка модели
print("Загружаем модель...")
model = AutoModel.from_pretrained("prostochel097/alpha3D-v1", trust_remote_code=True)
model.eval()

def generate_3d(red, green, blue, cube, sphere):
    # Собираем вектор из ползунков
    # Порядок в словаре был: red, green, blue, cube, sphere
    # Мы просто собираем их в список
    vec = [float(red), float(green), float(blue), float(cube), float(sphere)]
    tensor_in = torch.tensor([vec], dtype=torch.float32)
    
    # Генерация
    with torch.no_grad():
        out = model(tensor_in).numpy()[0]
    
    # Разбираем данные
    x = out[:, 0]
    y = out[:, 1]
    z = out[:, 2]
    
    # Цвета (0..1) -> (0..255) для Plotly
    rgb = np.clip(out[:, 3:], 0, 1)
    colors = [f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})' for r,g,b in rgb]

    # Рисуем интерактивный график Plotly
    fig = go.Figure(data=[go.Scatter3d(
        x=x, y=y, z=z,
        mode='markers',
        marker=dict(
            size=4,
            color=colors, # Индивидуальный цвет каждой точки
            opacity=0.8
        )
    )])
    
    # Настройки камеры и осей
    fig.update_layout(
        scene=dict(
            xaxis=dict(range=[-1.5, 1.5]),
            yaxis=dict(range=[-1.5, 1.5]),
            zaxis=dict(range=[-1.5, 1.5]),
            aspectmode='cube'
        ),
        margin=dict(r=0, l=0, b=0, t=0)
    )
    
    return fig

# 2. Интерфейс Gradio
with gr.Blocks() as demo:
    gr.Markdown("# 🧊 alpha3D-v1 Demo")
    gr.Markdown("Генерация 3D облака точек из вектора свойств. Покрутите ползунки, чтобы смешивать формы и цвета!")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 🎨 Colors")
            s_red = gr.Slider(0, 1, value=1, step=0.1, label="Red")
            s_green = gr.Slider(0, 1, value=0, step=0.1, label="Green")
            s_blue = gr.Slider(0, 1, value=0, step=0.1, label="Blue")
            
            gr.Markdown("### 📐 Shapes")
            s_cube = gr.Slider(0, 1, value=0, step=0.1, label="Cube")
            s_sphere = gr.Slider(0, 1, value=1, step=0.1, label="Sphere")
            
            btn = gr.Button("Generate 3D", variant="primary")
            
        with gr.Column():
            # Окно вывода 3D
            plot = gr.Plot(label="3D Output")

    # Связываем кнопку с функцией
    btn.click(
        fn=generate_3d, 
        inputs=[s_red, s_green, s_blue, s_cube, s_sphere], 
        outputs=plot
    )
    
    # Запускаем генерацию сразу при открытии
    demo.load(
        fn=generate_3d, 
        inputs=[s_red, s_green, s_blue, s_cube, s_sphere], 
        outputs=plot
    )

demo.launch()