from pathlib import Path
import gradio as gr
import numpy as np
import plotly.graph_objects as go

POINT_BUDGET = 1500
DATA_DIR = Path(__file__).parent / "data"

def load_data_files():
    data_arrays = []
    for i in range(5):
        file_path = DATA_DIR / f"{i}.npy"
        if file_path.exists():
            data_arrays.append(np.load(file_path, mmap_mode="r").reshape(-1))
    return data_arrays

DATA_SOURCES = load_data_files()
TOTAL_SAMPLES = len(DATA_SOURCES[0]) if DATA_SOURCES else 1_000_000


def create_figure(start_pt, window_size):

    end_pt = min(start_pt + window_size, TOTAL_SAMPLES)
    view_size = end_pt - start_pt

    dynamic_stride = max(1, view_size // POINT_BUDGET)
    indices = np.arange(start_pt, end_pt, dynamic_stride)

    fig = go.Figure()
    colors = ['#00E5FF', '#FF3D00', '#76FF03', '#D500F9', '#FFC400']

    for i, data in enumerate(DATA_SOURCES):
        fig.add_trace(go.Scattergl(
            x=indices,
            y=data[indices],
            mode="lines",
            name=f"Trace {i}",
            line=dict(width=1, color=colors[i]),
            hoverinfo="skip"
        ))

    fig.update_layout(
        template="plotly_dark",
        title=f"Stride: {dynamic_stride}x | Window: {start_pt:,} → {end_pt:,}",
        height=600
    )

    return fig


with gr.Blocks() as demo:

    gr.Markdown("### 🔬 ASCAD Trace Viewer (Stable HF Version)")

    start_slider = gr.Slider(
        minimum=0,
        maximum=TOTAL_SAMPLES - 1000,
        step=100,
        value=0,
        label="Start Index"
    )

    window_slider = gr.Slider(
        minimum=1000,
        maximum=50000,
        step=1000,
        value=20000,
        label="Window Size"
    )

    trace_plot = gr.Plot()

    start_slider.change(
        create_figure,
        inputs=[start_slider, window_slider],
        outputs=trace_plot
    )

    window_slider.change(
        create_figure,
        inputs=[start_slider, window_slider],
        outputs=trace_plot
    )

    # Initial load
    demo.load(
        create_figure,
        inputs=[start_slider, window_slider],
        outputs=trace_plot
    )


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