import gradio as gr
import numpy as np

def cascade_demo(query, N, r):
    scores = np.log(r**np.arange(int(N))) + query/100
    top5 = np.argsort(scores)[-5:][::-1]
    
    result = "TOP 5 DOCS: "
    for i in range(5):
        idx = int(top5[i])
        score = float(scores[idx])
        result = result + "Doc#" + str(idx) + ":" + str(round(score,2)) + " "
    
    metrics = "N=" + str(int(N)) + " r=" + str(round(r,3)) + " stable"
    return result, metrics

with gr.Blocks() as demo:
    gr.Markdown("# CascadeRAG Demo")
    with gr.Row():
        q = gr.Slider(0, 100, 50, label="Query")
        n = gr.Slider(50, 500, 100, label="Docs") 
        r = gr.Slider(0.8, 0.99, 0.92, label="r")
    
    out1 = gr.Textbox(lines=4, label="Results")
    out2 = gr.Textbox(lines=2, label="Metrics")
    
    gr.Button("Run").click(cascade_demo, [q,n,r], [out1,out2])

demo.launch()