# app.py
import json
import gradio as gr

# Load the merged knowledge base
with open("knowledge_base.json", encoding="utf-8") as f:
    kb = json.load(f)

symptom_to_icd    = kb["symptom_to_icd"]
icd_to_description = kb["icd_to_description"]

def map_symptoms(raw_input):
    terms = [t.strip().lower() for t in raw_input.split(",") if t.strip()]
    icd_counts = {}
    for term in terms:
        for code in symptom_to_icd.get(term, []):
            icd_counts[code] = icd_counts.get(code, 0) + 1
    if not icd_counts:
        return {"diagnoses": [], "confidences": []}
    total = sum(icd_counts.values())
    # sort codes by frequency descending
    sorted_items = sorted(icd_counts.items(), key=lambda x: x[1], reverse=True)
    diagnoses = []
    confidences = []
    for code, count in sorted_items:
        desc = icd_to_description.get(code, "Unknown")
        diagnoses.append(f"{code}: {desc}")
        confidences.append(round(count / total, 2))
    return {"diagnoses": diagnoses, "confidences": confidences}

# Use Blocks so that mcp_server=True is accepted
with gr.Blocks(mcp_server=True) as demo:
    gr.Markdown("## Symptom to ICD‐10 Code Lookup")
    inp = gr.Textbox(label="Enter symptoms (comma‐separated)")
    out = gr.JSON(label="Result")
    # Wire the submit event
    inp.submit(fn=map_symptoms, inputs=inp, outputs=out)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)