import gradio as gr
import numpy as np
from robot import TrajectoryPlanner
import logging
import plotly.graph_objects as go
from plotly.subplots import make_subplots
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
import time
import threading

# Initialize trajectory planner
try:
    planner = TrajectoryPlanner()
    PLANNER_AVAILABLE = True
except Exception as e:
    logging.error(f"Failed to initialize trajectory planner: {e}")
    PLANNER_AVAILABLE = False

animation_running = False
animation_thread = None

css = """
.control-panel .gr-slider { 
    margin-top: 4px !important;
    margin-bottom: 4px !important;
}
.control-panel .gr-form {
    gap: 6px !important;
}
.label-override {
    font-weight: 500 !important;
    font-size: 1.1em !important;
    color: #333 !important;
}
"""

def create_comparison_plots(plot_data):
    """Create comparison plots for the trajectories"""
    if plot_data is None:
        return None
    
    ours = plot_data['ours']
    toppra = plot_data['toppra']
    
    # Create subplots
    fig = make_subplots(
        rows=3, cols=2,
        subplot_titles=('Position', 'Velocity', 'Acceleration', 'Jerk', 'Position Error', 'Orientation Error'),
        vertical_spacing=0.12,
        horizontal_spacing=0.1
    )
        
    # Position
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['x'], name='Ours', 
                            line=dict(color='#6C8EBF', width=2), showlegend=True), row=1, col=1)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['x'], name='TOPPRA', 
                            line=dict(color='#B85450', width=2), showlegend=True), row=1, col=1)
    
    # Velocity
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['dx'], name='Ours', 
                            line=dict(color='#6C8EBF', width=2), showlegend=False), row=1, col=2)
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['dx_max'], name='Ours Limits', 
                            line=dict(color='#6C8EBF', width=1, dash='dash'), showlegend=True), row=1, col=2)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['dx'], name='TOPPRA', 
                            line=dict(color='#B85450', width=2), showlegend=False), row=1, col=2)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['ds_max'], name='TOPPRA Limits', 
                            line=dict(color='#B85450', width=1, dash='dash'), showlegend=True), row=1, col=2)
    
    # Acceleration
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['ddx'], name='Ours', 
                            line=dict(color='#6C8EBF', width=2), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['ddx_max'], name='Ours Limits', 
                            line=dict(color='#6C8EBF', width=1, dash='dash'), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['ddx_min'], name='Ours Limits', 
                            line=dict(color='#6C8EBF', width=1, dash='dash'), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['ddx'], name='TOPPRA', 
                            line=dict(color='#B85450', width=2), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['dds_max'], name='TOPPRA Limits', 
                            line=dict(color='#B85450', width=1, dash='dash'), showlegend=False), row=2, col=1)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['dds_min'], name='TOPPRA Limits', 
                            line=dict(color='#B85450', width=1, dash='dash'), showlegend=False), row=2, col=1)
    
    # Jerk
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['dddx'], name='Ours', 
                            line=dict(color='#6C8EBF', width=2), showlegend=False), row=2, col=2)
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['dddx_max'], name='Ours Limits', 
                            line=dict(color='#6C8EBF', width=1, dash='dash'), showlegend=False), row=2, col=2)
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['dddx_min'], name='Ours Limits', 
                            line=dict(color='#6C8EBF', width=1, dash='dash'), showlegend=False), row=2, col=2)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['dddx'], name='TOPPRA', 
                            line=dict(color='#B85450', width=2), showlegend=False), row=2, col=2)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['ddds_max'], name='TOPPRA Limits', 
                            line=dict(color='#B85450', width=1, dash='dash'), showlegend=False), row=2, col=2)  
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['ddds_min'], name='TOPPRA Limits', 
                            line=dict(color='#B85450', width=1, dash='dash'), showlegend=False), row=2, col=2)  
    
    # Position Error
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['e_pos'], name='Ours', 
                            line=dict(color='#6C8EBF', width=2), showlegend=False), row=3, col=1)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['e_pos'], name='TOPPRA', 
                            line=dict(color='#B85450', width=2), showlegend=False), row=3, col=1)
    
    # Orientation Error 
    fig.add_trace(go.Scatter(x=ours['t'], y=ours['e_rot'], name='Ours',
                            line=dict(color='#6C8EBF', width=2), showlegend=False), row=3, col=2)
    fig.add_trace(go.Scatter(x=toppra['t'], y=toppra['e_rot'], name='TOPPRA',
                            line=dict(color='#B85450', width=2), showlegend=False), row=3, col=2)
    
    
    # Update axes
    fig.update_xaxes(title_text="Time [s]", row=3, col=1)
    fig.update_xaxes(title_text="Time [s]", row=3, col=2)
    
    fig.update_yaxes(title_text="s [m]", row=1, col=1)
    fig.update_yaxes(title_text="ds [m/s]", row=1, col=2)
    fig.update_yaxes(title_text="dds [m/s²]", row=2, col=1)
    fig.update_yaxes(title_text="ddds [m/s³]", row=2, col=2)
    fig.update_yaxes(title_text="Position Error [mm]", row=3, col=1)
    fig.update_yaxes(title_text="Orientation Error [deg]", row=3, col=2)
    
    fig.update_layout(height=600, showlegend=True, legend=dict(x=0.5, y=1.1, orientation='h'))
    
    return fig

def animate_trajectory():
    """Background thread to animate the trajectory"""
    global animation_running
    
    anim_data = planner.get_animation_data()
    if anim_data is None:
        return
    
    t_max = anim_data['t_max']
    
    for i in range(10):
        if not animation_running:
            return
        t0 = time.time()
        
        while animation_running and (time.time() - t0) < t_max:
            t_current = time.time() - t0
            planner.update_animation(t_current)
            time.sleep(0.01)
            
    animation_running = False  
    

def generate_and_compute(traj_length, capacity_scale, progress=gr.Progress()):
    """Generate trajectory and compute both algorithms"""
    if not PLANNER_AVAILABLE:
        return "Planner not available", None, None
    
    progress(0, desc="Generating trajectory...")
    
    # try:
    result = planner.generate_trajectory(traj_length, capacity_scale, progress)
    
    info = f"""
    ✅ **Trajectory Generated Successfully**
    
    - **Waypoints**: {result['waypoints']}
    - **Trajectory Length**: {traj_length:.2f} m
    - **Capacity Scale**: {capacity_scale:.2f}
    - **TOPPRA Duration**: {result['toppra_duration']:.3f} s
    - **Ours Duration**: {result['ours_duration']:.3f} s
    - **Speedup**: {(result['toppra_duration'] / result['ours_duration']):.2f}x
    """
    
    progress(0.8, desc="Creating plots...")
    
    plot_data = planner.get_plot_data()
    plots = create_comparison_plots(plot_data)
    
    progress(1.0, desc="Done!")
    
    return info, plots, gr.update(interactive=True)
    
    #except Exception as e:
    #    logging.error(f"Error generating trajectory: {e}")
    #    return f"❌ Error: {str(e)}", None, gr.update(interactive=False)

def start_animation():
    """Start the animation"""
    global animation_running, animation_thread
    
    if not PLANNER_AVAILABLE or planner.data is None:
        return gr.update(value="⚠️ Generate trajectory first")
    
    animation_running = True
    animation_thread = threading.Thread(target=animate_trajectory, daemon=True)
    animation_thread.start()
    
    return gr.update(value="⏸️ Stop Animation", variant="stop")

def stop_animation():
    """Stop the animation"""
    global animation_running
    animation_running = False
    
    return gr.update(value="▶️ Play Animation", variant="primary")

def toggle_animation(btn_state):
    """Toggle animation on/off"""
    if "Stop" in btn_state:
        return stop_animation()
    else:
        return start_animation()

with gr.Blocks(css=css, title="Capacity-Aware Trajectory Planning") as demo:
    
    gr.Markdown("""
    # 🤖 Capacity-Aware Trajectory Planning Comparison
    
    📄 This demo implements the paper: **Online approach to near time-optimal task-space trajectory planning**
    by Antun Skuric, Nicolas Torres Alberto, Lucas Joseph, Vincent Padois, David Daney
    Link: https://inria.hal.science/hal-04576076 

    
    
    Compare **capacity-aware real-time trajectory planning** vs **TOPPRA** for the Panda robot.
    
    **Left Robot**: Our approach | **Right Robot**: TOPPRA
    """)

    with gr.Row(equal_height=True):
        # LEFT: MeshCat viewer
        with gr.Column(scale=2, min_width=500):
            if PLANNER_AVAILABLE:
                viewer = gr.HTML(planner.iframe(height=600))
            else:
                viewer = gr.HTML("<div style='height:600px; background:#f0f0f0; display:flex; align-items:center; justify-content:center;'><h2>Planner Not Available</h2></div>")

        # RIGHT: Controls
        with gr.Column(scale=1, min_width=400, elem_classes="control-panel"):
            gr.Markdown("## 🎛️ Trajectory Configuration")
            
            traj_length = gr.Slider(
                0.1, 1.0, value=0.8, step=0.05,
                label="Trajectory Length (m)",
                info="Distance between start and end points"
            )
            
            capacity_scale = gr.Slider(
                0.1, 1.0, value=0.5, step=0.05,
                label="Robot Capacity Scale",
                info="Fraction of maximum robot capacity (1.0 = full capacity)"
            )
            
            generate_btn = gr.Button("🎲 Generate Random Trajectory", variant="primary", size="lg")
            
            info_box = gr.Markdown("Click 'Generate' to create a random trajectory")
            
            gr.Markdown("## 🎬 Animation Control")
            
            play_btn = gr.Button("▶️ Play Animation", variant="primary", interactive=False)
            
    with gr.Row():
        with gr.Column():
            gr.Markdown("## 📊 Trajectory Comparison")
            plots = gr.Plot(label="Comparison Plots")
    
    # Callbacks
    generate_btn.click(
        generate_and_compute,
        inputs=[traj_length, capacity_scale],
        outputs=[info_box, plots, play_btn]
    )
    
    play_btn.click(
        toggle_animation,
        inputs=[play_btn],
        outputs=[play_btn]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=8501)