import os
os.environ['MUJOCO_GL'] = 'osmesa'
os.environ['PYOPENGL_PLATFORM'] = 'osmesa'

import streamlit as st
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import torch as th
from ae import *
from geometry import *
try:
    from clustering_robots import *
except AttributeError as e:
    st.error("❌ failed to load library!")
    st.stop()
from torch_chain import *
from transforms import *

st.markdown("""
    <style>
        #MainMenu {visibility: hidden;}
        footer {visibility: hidden;}
        header {visibility: hidden;}
        .block-container {padding-top: 1rem; padding-bottom: 0rem;}
    </style>
""", unsafe_allow_html=True)
st.set_page_config(layout="wide") # 넓은 화면 모드
st.title("🤖 LEGO Latent Space Explorer")
st.markdown("""
This app explores **Latent Space** and decodes kinematic structure in real time.
Left: Current Design Space, Right: Decoded Kinematic Structure
""")

@st.cache_resource
def init_robot_system():
    ae, config = load_saved_model(device="cpu")
    robot_chars = RobotCharacteristics()
    njnt_data = np.load('./src/njnt.npy')
    robot_chars.max_joints_per_group = njnt_data
    return ae, robot_chars




ae, robot_chars = init_robot_system()

@st.cache_data(show_spinner=False)
def load_motion_data(level):
    """
    load meshgrid per level
    """
    # check whether it is valid path
    x_1 = np.load(f'./src/{level}_X1.npy')
    x_2 = np.load(f'./src/{level}_X2.npy')
    z_1 = np.load(f'./src/{level}_Z1.npy')
    return x_1, x_2, z_1

with st.expander("🎛️ Control Latent Vector (Click to Open/Close)", expanded=True):
    # 1. radio button
    selected_level = st.radio(
        "Select Motion Level",
        options=["easy", "middle", "hard"], 
        index=0,
        horizontal=True
    )
    
    with st.spinner(f"Loading {selected_level} data..."):
        x_1, x_2, z_1 = load_motion_data(selected_level)

    st.write("---")
    
    # slider
    c_s1, c_s2 = st.columns(2)
    l1 = c_s1.slider("Latent X (l1)", float(x_1.min()), float(x_1.max()), 0.0)
    l2 = c_s2.slider("Latent Y (l2)", float(x_2.min()), float(x_2.max()), 0.0)

h1, h2 = st.columns([1, 1])
h1.subheader("Latent Space (3D Mesh)")
h2.subheader("Kinematic Structure")


col1, col2 = st.columns([1, 1], vertical_alignment="center")

with col1:
    
    fig = go.Figure()
    
    fig.add_trace(go.Scatter3d(
        x=x_1.flatten(), y=x_2.flatten(), z=z_1.flatten(),
        mode='markers',
        marker=dict(size=2, color=z_1.flatten(), colorscale='Viridis', opacity=0.5),
        name='Latent Manifold'
    ))
    
    current_z = np.interp(l1, x_1, z_1) if x_1.ndim==1 else 0
    fig.add_trace(go.Scatter3d(
        x=[l1], y=[l2], z=[current_z],
        mode='markers',
        marker=dict(size=8, color='red', symbol='diamond'),
        name='Current Point'
    ))

    # layout
    fig.update_layout(
        height=400, 
        margin=dict(l=0, r=0, b=0, t=0),
        legend=dict(
            yanchor="top", y=0.9,
            xanchor="center", x=0.5,
            orientation="h",
            bgcolor="rgba(0,0,0,0)"
        ),
        scene=dict(
            xaxis_title='X1', yaxis_title='X2', zaxis_title='Z1',
            aspectmode='cube',
            xaxis=dict(showbackground=False),
            yaxis=dict(showbackground=False),
            zaxis=dict(showbackground=False),
        )
    )
    st.plotly_chart(fig, use_container_width=True)

with col2:
    # decode logic
    opt = np.array([[l1, l2]])
    
    with st.spinner("Generating..."):
        try:
            decode = ae.decode(th.tensor(opt, device=next(ae.parameters()).device, dtype=th.float32)).detach().numpy()
            full_decode = robot_chars.build_full_interp_from_symmetric(decode, threshold=0.0)
            data_to_visualize = robot_chars.convert_flat_interp_to_dict(full_decode)
            
            # image
            result_img = robot_chars.visualize_data_to_classify(data_to_visualize, vis=True)
            
            if result_img:
                # fit to center
                st.image(result_img, use_container_width=True)
                
        except Exception as e:
            st.error(f"Visualization Error: {e}")