import gradio as gr
from data_loader import load_data_from_file
from visualization import create_3d_visualization
import os

LANDMARKS = ["SHOULDER", "ELBOW", "WRIST", "INDEX", "THUMB"]

with gr.Blocks() as demo:
    gr.Markdown("# **Visualisation 3D Interactive du Mouvement**")

    with gr.Row():
        with gr.Column(scale=1, min_width=300):
            file_input = gr.File(label="Charger un fichier CSV")
            file_status = gr.Textbox(label="Statut du fichier", interactive=False)

            landmark_checkboxes = {}
            with gr.Accordion("Visibilité des repères", open=True):
                for lm in LANDMARKS:
                    landmark_checkboxes[lm] = gr.Checkbox(value=True, label=lm)

            with gr.Accordion("Options d'affichage", open=True):
                show_trajectories_checkbox = gr.Checkbox(value=True, label="Afficher les trajectoires")
                show_segments_checkbox = gr.Checkbox(value=True, label="Afficher les segments")

            frame_slider = gr.Slider(minimum=0, maximum=0, step=1, value=0, label="Frame")
            frame_label = gr.Textbox(label="Frame actuelle", interactive=False)

        with gr.Column(scale=4):
            output_plot = gr.Plot(label="Visualisation 3D")

    df_state = gr.State(None)
    n_frames_state = gr.State(0)
    current_frame_state = gr.State(0)

    def on_file_upload(file_obj):
        if file_obj is None:
            return None, 0, "Aucun fichier chargé", 0, gr.Slider(maximum=0, value=0), "Frame: 0/0", None

        df, n_frames = load_data_from_file(file_obj)
        if df is None:
            return None, 0, "Erreur de chargement", 0, gr.Slider(maximum=0, value=0), "Frame: 0/0", None

        fig = create_3d_visualization(df, 0)
        
        return (df, n_frames, f"Fichier: {os.path.basename(file_obj.name)}", 
                0, gr.Slider(maximum=n_frames-1, value=0), 
                f"Frame: 0/{n_frames-1}", fig)

    file_input.upload(
        on_file_upload,
        inputs=[file_input],
        outputs=[df_state, n_frames_state, file_status, current_frame_state, 
                frame_slider, frame_label, output_plot]
    )

    def update_visualization_and_frame_label(frame, df, n_frames, show_traj, show_seg, *lm_vis_values):
        if df is None:
            return None, "Frame: 0/0"

        lm_visibility = {lm: val for lm, val in zip(LANDMARKS, lm_vis_values)}
        fig = create_3d_visualization(df, int(frame), show_traj, show_seg, lm_visibility)

        return fig, f"Frame: {int(frame)}/{n_frames-1}"

    frame_slider.change(
        update_visualization_and_frame_label,
        inputs=[frame_slider, df_state, n_frames_state, show_trajectories_checkbox, show_segments_checkbox] + [landmark_checkboxes[lm] for lm in LANDMARKS],
        outputs=[output_plot, frame_label]
    )

    for cb in [show_trajectories_checkbox, show_segments_checkbox] + [landmark_checkboxes[lm] for lm in LANDMARKS]:
        cb.change(
            update_visualization_and_frame_label,
            inputs=[frame_slider, df_state, n_frames_state, show_trajectories_checkbox, show_segments_checkbox] + [landmark_checkboxes[lm] for lm in LANDMARKS],
            outputs=[output_plot, frame_label]
        )

demo.launch()