import gradio as gr
from scatter_plot_snr import plot_mass_vs_distance_or_redshift

def plot_interactive(snr_list, alpha, y_axis, x_axis, colorbar_var):
    fig = plot_mass_vs_distance_or_redshift(
        snrs=snr_list,
        alpha=alpha,
        y_axis=y_axis,
        x_axis=x_axis,
        colorbar_var=colorbar_var
    )
    import io
    import PIL.Image
    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=300)
    buf.seek(0)
    img = PIL.Image.open(buf)
    return img

with gr.Blocks() as demo:
    gr.Markdown("""
        ## EMRI Search Mass vs Distance/Redshift Interactive Plot

        This visualization shows the sources detected by the search pipeline presented in [Ab uno disce omnes: Single-harmonic search for extreme mass-ratio inspirals](https://arxiv.org/abs/2510.20891).
        On the x-axis, you can choose to plot either the primary or secondary mass of the detected sources. 
        On the y-axis, you can choose to plot either the redshift or luminosity distance. 
        The color of the points corresponds to a variable of your choice (initial eccentricity, final eccentricity, primary mass, or secondary mass).
        Since the injection was performed in the detector frame, the masses are plotted in the source frame by dividing by (1+z).
        
        For a quick, simplified EMRI search, open and run [QuickStartEMRIsearch.ipynb](examples/QuickStartEMRIsearch.ipynb). The notebook walks through the pipeline and produces example outputs, including:

        If you use this code, please cite the [paper](https://ui.adsabs.harvard.edu/abs/2026PhRvD.113b4061S)
        """)
    snr_choices = [20, 25, 30, 35, 40]
    with gr.Row():
        snr_input = gr.CheckboxGroup(snr_choices, label="Select SNRs", value=[30])
        alpha_input = gr.Number(label="False Alarm Rate (alpha)", value=1e-4)
    with gr.Row():
        y_axis_input = gr.Radio(["Redshift", "Luminosity Distance"], label="Y Axis", value="Redshift")
        x_axis_input = gr.Radio(["Primary Mass", "Secondary Mass"], label="X Axis", value="Primary Mass")

    colorbar_input = gr.Dropdown([r"Final eccentricity", r"Initial eccentricity", 
                                  r"Primary mass", r"Secondary mass"], label="Colorbar Variable", value=r"Final eccentricity")
    plot_btn = gr.Button("Click here to plot")
    plot_output = gr.Image(type="pil")

    plot_btn.click(
        plot_interactive,
        inputs=[snr_input, alpha_input, y_axis_input, x_axis_input, colorbar_input],
        outputs=plot_output
    )

demo.launch()