app.py

import os
import sys
import tempfile
import zipfile
import shutil
from pathlib import Path
from datetime import datetime

os.environ["MPLBACKEND"] = "Agg"
os.environ["QT_QPA_PLATFORM"] = "offscreen"

APP_DIR = Path(__file__).resolve().parent
SDK_PARENT = APP_DIR / "accerionsdk" / "tools" / "python"
sys.path.insert(0, str(SDK_PARENT))
sys.path.insert(0, str(APP_DIR))

import gradio as gr
from sdk.data.data_container import DataContainer
from sdk.utils.configuration import Configuration

# All arams-based plotting tools
from sdk.plotting.plot_distance_between_corrections_histogram import plot_distance_between_corrections_histogram
from sdk.plotting.plot_distance_between_corrections_map import plot_distance_between_corrections_map
from sdk.plotting.plot_distance_between_signatures_histogram import plot_distance_between_signatures_histogram
from sdk.plotting.plot_distance_between_signatures_map import plot_distance_between_signatures_map
from sdk.plotting.plot_external_reference_timeseries import plot_external_reference_timeseries
from sdk.plotting.plot_heading_corrections_histogram import plot_heading_corrections_histogram
from sdk.plotting.plot_heading_corrections_map import plot_heading_corrections_map
from sdk.plotting.plot_line_following_position_errors_histogram import plot_line_following_position_errors_histogram
from sdk.plotting.plot_line_following_position_errors_map import plot_line_following_position_errors_map
from sdk.plotting.plot_map_visualization import plot_map_visualization
from sdk.plotting.plot_mapping_direction import plot_mapping_direction
from sdk.plotting.plot_position_corrections_histogram import plot_position_corrections_histogram
from sdk.plotting.plot_position_corrections_map import plot_position_corrections_map
from sdk.plotting.plot_sensor_poses import plot_sensor_poses
from sdk.plotting.plot_residual_position_errors_g2o import plot_residual_position_errors_g2o
from sdk.plotting.plot_residual_heading_errors_g2o import plot_residual_heading_errors_g2o


PLOT_CHOICES_MAP = {
    # Floor-map / signature position plots
    "Map visualization": [plot_map_visualization, plot_mapping_direction],
    "Distance between signatures ": [plot_distance_between_signatures_map, plot_distance_between_signatures_histogram],

    # G2O / optimization-related plots
    "Residual loop closure errors (g2o)": [plot_residual_position_errors_g2o, plot_residual_heading_errors_g2o],
}

PLOT_CHOICES_LOCALIZATION = {
    # Arams-based plots
    "Distance between corrections": [plot_distance_between_corrections_map, plot_distance_between_corrections_histogram],
    "Position corrections": [plot_position_corrections_map, plot_position_corrections_histogram],
    "Heading corrections": [plot_heading_corrections_map, plot_heading_corrections_histogram],
    "External reference": [plot_external_reference_timeseries, plot_sensor_poses],
    "Line-following position errors": [plot_line_following_position_errors_map, plot_line_following_position_errors_histogram],
}

PLOT_CHOICES_LINE_FOLLOWING = {
    # Line-following plots
}


def find_arams_dir(root: Path):
    hits = [p for p in root.rglob("*") if p.is_dir() and p.name.lower() == "arams_user_logs"]
    return sorted(hits, key=lambda p: len(p.parts))[-1] if hits else None


def run_pipeline(
    zip_file,
    map_file,
    g2o_file,
    overlay_map,
    overlay_input_poses,
    visualize_outliers,
    visualize_first_drift,
    plot_time_interval,
    skip_initial_drift,
    selected_plots_map,
    selected_plots_loc,
):
    if zip_file is None:
        return [], None, "❌ No log zip selected."

    workdir = Path(tempfile.mkdtemp(prefix="plots_"))
    in_dir = workdir / "input"
    out_dir = workdir / "output"
    in_dir.mkdir(parents=True, exist_ok=True)
    out_dir.mkdir(parents=True, exist_ok=True)

    ### 1) unzip logs
    try:
        with zipfile.ZipFile(zip_file, "r") as zf:
            zf.extractall(in_dir)
    except zipfile.BadZipFile:
        shutil.rmtree(workdir, ignore_errors=True)
        return [], None, "❌ Invalid zip file."

    ### 2) find arams_user_logs directory
    arams_dir = find_arams_dir(in_dir)
    if not arams_dir:
        try:
            top_list = "\n- ".join(str(p.relative_to(in_dir)) for p in sorted(in_dir.iterdir())[:20])
        except Exception:
            top_list = "(unavailable)"
        return [], None, f"❌ 'arams_user_logs' folder not found.\nTop-level after unzip:\n- {top_list}"

    ### 3) prepare configuration
    try:
        conf = Configuration.load_configuration(out_dir)

        # Map plotting options
        conf[Configuration.Name.OVERLAY_MAP] = bool(overlay_map)
        conf[Configuration.Name.OVERLAY_INPUT_POSES] = bool(overlay_input_poses)
        conf[Configuration.Name.VISUALIZE_OUTLIERS] = bool(visualize_outliers)
        conf[Configuration.Name.VISUALIZE_FIRST_DRIFT_CORRECTION] = bool(visualize_first_drift)

        # Filters
        conf[Configuration.Name.PLOT_TIME_INTERVAL] = int(plot_time_interval*60*10**9)  # Convert to nanoseconds

        # Drift / clusters
        conf[Configuration.Name.SKIP_INITIAL_DRIFT_CORRECTIONS] = int(skip_initial_drift or 0)

        # TODO: Change the metrics values based on inputs from 'Metrics' input tab
        # from sdk.metrics.reader import Reader as MetricsReader
        # from sdk.metrics.metrics_data import MetricsData
        # from sdk.metrics.writer import Writer
        # conf[Configuration.Name.METRICS] = MetricsReader.read_default()

        ### 4) Gather input files
        input_files = [p for p in arams_dir.rglob("*") if p.is_file()]

        # if map_file is added: copy into working dir & append
        if map_file is not None:
            map_path = Path(map_file.name)
            extra_path = workdir / map_path.name
            shutil.copy(map_file.name, extra_path)
            input_files.append(extra_path)

        # if g2o file is added: copy into working dir & append
        if g2o_file is not None:
            g2o_path = Path(g2o_file.name)
            extra_path = workdir / g2o_path.name
            shutil.copy(g2o_file.name, extra_path)
            input_files.append(extra_path)

        ### 5) run selected plots
        data = DataContainer.create(input_files, conf)
        for s in selected_plots_map:
            scripts = PLOT_CHOICES_MAP[s]
            for script in scripts:
                script(data, out_dir, verbose=True, conf=conf)

        for s in selected_plots_loc:
            scripts = PLOT_CHOICES_LOCALIZATION[s]
            for script in scripts:
                script(data, out_dir, verbose=True, conf=conf)

    except Exception as e:
        return [], None, f"❌ Error while preparing/running plots:\n{e}"

    ### 6) gather preview images
    images = []
    for p in sorted(out_dir.rglob("*")):
        if p.is_file() and p.suffix.lower() in {".png"}:
            images.append(str(p))

    ### 7) zip outputs
    zip_name = f"plots_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.zip"
    zip_path = workdir / zip_name
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for p in out_dir.rglob("*"):
            if p.is_file():
                zf.write(p, arcname=str(p.relative_to(out_dir)))

    status = f"✅ plots finished. {len(images)} previewable files found. Download: {zip_name}"
    return images, str(zip_path), status


# -------- UI --------
with gr.Blocks(title="Visualize Triton data", delete_cache={86400, 86400}) as demo:
    gr.Markdown("### Upload your log `.zip` and optional `.db` map file, set parameters, and click **Run**.")

    with gr.Row():
        file_in = gr.File(label="Select logs (.zip)", file_types=[".zip"])
        map_file_in = gr.File(label="Optional: map file (.db)", file_types=[".db"])
        g2o_file_in = gr.File(label="Optional: loop closure file (.g2o)", file_types=[".g2o"])

    with gr.Tab("Plot selection"):
        with gr.Row():
            selected_plots_map = gr.CheckboxGroup(choices=list(PLOT_CHOICES_MAP.keys()),
                                                  label="Map")
        with gr.Row():
            selected_plots_loc = gr.CheckboxGroup(choices=list(PLOT_CHOICES_LOCALIZATION.keys()),
                                                  label="Localization")

    with gr.Tab("Plot parameters"):
        gr.Markdown("**Time interval filter**")
        with gr.Row():
            plot_time_interval = gr.Number(value=5, label="Plot last N minutes (use '-1' for no filter)")

        gr.Markdown("**Drift correction filter**")
        with gr.Row():
            skip_initial_drift = gr.Number(value=0, precision=0, label="Skip initial drift corrections")

        gr.Markdown("**Other visualization options**")
        with gr.Row():
            overlay_map = gr.Checkbox(value=True, label="Overlay Triton map")
            overlay_input_poses = gr.Checkbox(value=False, label="Overlay external reference poses")
            visualize_outliers = gr.Checkbox(value=True, label="Secondary colorbar for large values")
            visualize_first_drift = gr.Checkbox(value=True, label="Visualize first drift correction")

    # with gr.Tab("Metrics"):
    #     gr.Markdown("**Loop closure errors - Position**")
    #     with gr.Row():
    #         max_lc_pos_error = gr.Number(value=.05, precision=2, label="Maximum desired value [m]")
    #         avg_lc_pos_error = gr.Number(value=.01, precision=2 ,label="Average desired value [m]")
    #
    #     gr.Markdown("**Loop closure errors - Heading**")
    #     with gr.Row():
    #         max_lc_heading_error = gr.Number(value=3.0, precision=1, label="Maximum desired value [deg]")
    #         avg_lc_heading_error = gr.Number(value=0.6, precision=1, label="Average desired value [deg]")
    #
    #     gr.Markdown("**Distance between signatures**")
    #     with gr.Row():
    #         max_sig_dist = gr.Number(value=0.075, precision=3, label="Maximum desired value [m]")
    #         min_sig_dist = gr.Number(value=0.0025, precision=4, label="Minimum desired value [m]")
    #         avg_sig_dist = gr.Number(value=0.020, precision=3, label="Average desired value [m]")
    #
    #     gr.Markdown("**Distance between consecutive corrections**")
    #     with gr.Row():
    #         max_dc_dist = gr.Number(value=0.75, precision=2, label="Maximum desired value [m]")
    #
    #     gr.Markdown("**Drift correction error - Position**")
    #     with gr.Row():
    #         max_dc_pos_error = gr.Number(value=0.05, precision=2, label="Maximum desired value [m]")
    #
    #     gr.Markdown("**Drift correction error - Heading**")
    #     with gr.Row():
    #         max_dc_heading_error = gr.Number(value=6., precision=1, label="Maximum desired value [deg]")
    #
    #     gr.Markdown("**Line-following error - Position**")
    #     with gr.Row():
    #         max_lf_pos_error = gr.Number(value=0.05, precision=2, label="Maximum desired value [m]")

    # Add 'Run' button that will run_pipeline when clicked
    run_btn = gr.Button("Run", variant="primary")

    with gr.Tabs():
        with gr.Tab("Preview"):
            gallery = gr.Gallery(label="Generated plots", columns=5, height=800)
        with gr.Tab("Download"):
            zip_out = gr.File(label="Download plots (.zip)")

    status = gr.Markdown()

    run_btn.click(
        run_pipeline,
        inputs=[
            file_in,
            map_file_in,
            g2o_file_in,
            overlay_map,
            overlay_input_poses,
            visualize_outliers,
            visualize_first_drift,
            plot_time_interval,
            skip_initial_drift,
            selected_plots_map,
            selected_plots_loc,
        ],
        outputs=[gallery, zip_out, status],
    )

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