"""
Gradio interface for SpikeCascade. Drop a .npy file containing a
(n_channels, n_samples) voltage array, and the Space returns the
sorted units, the calibration report, and the activity heatmap.
"""

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt

from modeling_spikecascade import SpikeCascadeForSorting


MODEL = SpikeCascadeForSorting.from_pretrained("your-org/SpikeCascade")


def sort(voltage_file):
    voltage = np.load(voltage_file.name).astype(np.float32)
    report = MODEL.infer(voltage)

    # Raster plot
    fig, ax = plt.subplots(figsize=(8, 3))
    for uid, times in report.units.items():
        ax.scatter(times, [uid] * len(times), s=4, marker="|")
    ax.set_xlabel("time (ms)")
    ax.set_ylabel("unit")
    ax.set_title(
        f"activity={report.activity_fraction:.3f}  "
        f"stage3={report.stage3_fraction:.3f}  "
        f"latency={report.latency_ms:.2f} ms  "
        f"corr={report.correlation:.3f}"
    )

    return (
        fig,
        report.activity_fraction,
        report.stage3_fraction,
        report.latency_ms,
        report.correlation,
    )


demo = gr.Interface(
    fn=sort,
    inputs=gr.File(label="voltage (.npy, shape = n_channels × n_samples)"),
    outputs=[
        gr.Plot(label="Spike raster"),
        gr.Number(label="activity_fraction"),
        gr.Number(label="stage3_fraction"),
        gr.Number(label="latency_ms"),
        gr.Number(label="correlation"),
    ],
    title="SpikeCascade — real-time Neuropixels spike sorting",
    description=(
        "Three-stage fidelity cascade. The calibration report is printed "
        "with every sort: activity_fraction, stage3_fraction, latency_ms, "
        "and correlation (NaN if the input contains zero spikes)."
    ),
)

demo.launch()