import json
from pathlib import Path
import gradio as gr
import plotly.graph_objects as go
import plotly.io as pio

DATA_DIR = Path(__file__).parent / "data"
video_path = DATA_DIR / "video.mp4"

METRIC_LABELS = {
    "x_cm": "X (cm)",
    "y_cm": "Y (cm)",
    "z_cm": "Z (cm)",
    "yaw_deg": "Yaw (°)",
    "pitch_deg": "Pitch (°)",
    "roll_deg": "Roll (°)",
}

PLOT_GRID = [
    ["x_cm", "y_cm", "z_cm"],
    ["yaw_deg", "pitch_deg", "roll_deg"],
]

PLOT_ORDER = [metric for row in PLOT_GRID for metric in row]

CUSTOM_CSS = """
:root, .gradio-container, body {
    background-color: #050a18 !important;
    color: #f8fafc !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif;
}
.side-panel {
    background: #0f172a;
    padding: 20px;
    border-radius: 18px;
    border: 1px solid #1f2b47;
    min-height: 100%;
}
.stats-card ul {
    list-style: none;
    padding: 0;
    margin: 0;
    font-size: 0.92rem;
}
.stats-card li {
    margin-bottom: 10px;
    color: #e2e8f0;
}
.stats-card span {
    display: inline-block;
    margin-right: 6px;
    color: #7dd3fc;
}
.main-panel {
    padding-top: 8px;
}
.video-card {
    background: #0f172a;
    border: 1px solid #1f2b47;
    border-radius: 18px;
    padding: 18px 20px;
    margin-top: 18px;
}
.video-title {
    font-size: 0.78rem;
    text-transform: uppercase;
    letter-spacing: 0.18em;
    color: #94a3b8;
    margin-bottom: 8px;
}
.video-panel video {
    border-radius: 12px;
    border: 1px solid #1f2b47;
    background: #030712;
}
.plots-wrap {
    margin-top: 18px;
}
.plots-wrap .gr-row {
    gap: 16px;
}
.plot-html {
    background: #111a2c;
    border-radius: 12px;
    padding: 10px;
    border: 1px solid #1f2b47;
    min-height: 320px;
}
.plot-html iframe {
    width: 100%;
    height: 300px;
    border: none;
}
"""

def load_data():
    metadata = {}
    end_effector = {}
    try:
        with open(DATA_DIR / "metadata.json", 'r') as f:
            metadata = json.load(f)
    except:
        pass
    try:
        with open(DATA_DIR / "end_effector.json", 'r') as f:
            end_effector = json.load(f)
    except:
        pass
    return metadata, end_effector

def build_plot_html(metadata, end_effector, hand, metric):
    try:
        fps = metadata.get('fps', 60)
        if not end_effector:
            return "<p>No data</p>"
        
        frame_keys = sorted([int(k) for k in end_effector.keys() if str(k).isdigit()])
        if not frame_keys:
            return "<p>No frames</p>"
        
        times = [i/fps for i in frame_keys]
        values = []
        
        for k in frame_keys:
            ee = end_effector.get(str(k), {}) or {}
            hd = ee.get(hand + "_hand")
            if hd and isinstance(hd, dict):
                pose = hd.get('pose_6dof', [])
                if len(pose) >= 6:
                    if metric == "x_cm":
                        values.append(pose[0] * 100)
                    elif metric == "y_cm":
                        values.append(pose[1] * 100)
                    elif metric == "z_cm":
                        values.append(pose[2] * 100)
                    elif metric == "roll_deg":
                        values.append(pose[3] * 57.3)
                    elif metric == "pitch_deg":
                        values.append(pose[4] * 57.3)
                    elif metric == "yaw_deg":
                        values.append(pose[5] * 57.3)
                else:
                    values.append(None)
            else:
                values.append(None)
        
        valid_times = [t for t, v in zip(times, values) if v is not None]
        valid_values = [v for v in values if v is not None]
        
        if not valid_times or not valid_values:
            return "<p>No valid data</p>"
        
        fig = go.Figure()
        fig.add_trace(go.Scatter(
            x=valid_times, 
            y=valid_values, 
            mode="lines",
            name=f"{hand} {metric}",
            line=dict(width=2)
        ))
        fig.update_layout(
            template="plotly_dark", 
            height=250, 
            margin=dict(l=20, r=20, t=30, b=20),
            xaxis_title="Time (s)",
            yaxis_title=METRIC_LABELS[metric]
        )
        return pio.to_html(fig, include_plotlyjs="cdn", full_html=False)
    except Exception as e:
        return f"<p>Error: {str(e)}</p>"

def build_interface():
    metadata, end_effector = load_data()
    total_frames = len(metadata.get('poses', []))
    fps = metadata.get('fps', 60)
    
    # Build plots for left and right hands
    left_figs = {metric: build_plot_html(metadata, end_effector, "left", metric) for metric in PLOT_ORDER}
    right_figs = {metric: build_plot_html(metadata, end_effector, "right", metric) for metric in PLOT_ORDER}
    
    stats_html = f"""
<div class="stats-card">
  <ul>
    <li><span>Number of samples/frames:</span> {total_frames:,}</li>
    <li><span>Frames per second:</span> {fps:.1f}</li>
  </ul>
</div>
"""
    
    theme = gr.themes.Soft(
        primary_hue="cyan", secondary_hue="blue", neutral_hue="slate"
    ).set(
        body_background_fill="#0c1424",
        body_text_color="#f8fafc",
        block_background_fill="#111a2c",
        block_title_text_color="#f8fafc",
        input_background_fill="#151f33",
        border_color_primary="#1f2b47",
        shadow_drop="none",
    )
    
    with gr.Blocks(theme=theme, css=CUSTOM_CSS) as demo:
        gr.Markdown("# 🤖 Dynamic Intelligence - Human Demo Visualizer")
        gr.Markdown("Egocentric hand tracking dataset for humanoid robot training.")
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=1, min_width=260, elem_classes=["side-panel"]):
                gr.HTML(stats_html)
                gr.HTML('<div class="episodes-title">Hands</div>')
                hand_radio = gr.Radio(
                    choices=["Left Hand", "Right Hand"],
                    value="Left Hand",
                    label="Hand Selection",
                )
            with gr.Column(scale=2, min_width=640, elem_classes=["main-panel"]):
                with gr.Column(elem_classes=["video-card"]):
                    gr.HTML('<div class="video-title">RGB Video</div>')
                    video = gr.Video(
                        height=360,
                        value=str(video_path) if video_path.exists() else None,
                        elem_classes=["video-panel"],
                        show_label=False,
                    )
        
        plot_outputs = []
        gr.Markdown("### Hand Trajectories")
        with gr.Column(elem_classes=["plots-wrap"]):
            for row in PLOT_GRID:
                with gr.Row():
                    for metric in row:
                        plot = gr.HTML(value=left_figs[metric], elem_classes=["plot-html"])
                        plot_outputs.append(plot)
        
        def update_plots(hand_choice):
            if hand_choice == "Left Hand":
                return [left_figs[metric] for metric in PLOT_ORDER]
            else:
                return [right_figs[metric] for metric in PLOT_ORDER]
        
        hand_radio.change(fn=update_plots, inputs=hand_radio, outputs=plot_outputs)
    
    return demo

# For Gradio Spaces
try:
    demo = build_interface()
except Exception as e:
    print(f"Error: {e}")
    import traceback
    traceback.print_exc()
    with gr.Blocks() as demo:
        gr.Markdown("# Error")
        gr.Markdown(str(e))

if __name__ == "__main__":
    demo.launch()