import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import soundfile as sf
import io

# ---------------------------------------------------------
# Utility Functions
# ---------------------------------------------------------

def process_image(img):
    if img is None:
        return "No image uploaded", None

    # Placeholder heatmap
    plt.figure(figsize=(4, 4))
    plt.imshow(np.random.rand(100, 100), cmap="jet")
    plt.title("AI Heatmap (Placeholder)")
    plt.axis("off")

    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    buf.seek(0)

    return "Image processed (placeholder inference).", buf


def process_audio(audio_file):
    if audio_file is None:
        return "No audio uploaded."

    # Read audio
    data, sample_rate = sf.read(audio_file)
    duration = len(data) / sample_rate

    return f"Cough audio received — Duration: {duration:.2f} sec (placeholder inference)"


def process_tabular(symptoms, age, smoker):
    if age is None:
        return "Please enter an age."

    risk = 0

    if symptoms:
        if "Persistent Cough" in symptoms:
            risk += 20
        if "Breathlessness" in symptoms:
            risk += 20
        if "Fever" in symptoms:
            risk += 15
        if "Fatigue" in symptoms:
            risk += 10

    if smoker:
        risk += 30

    risk += int(age) * 0.4

    return f"Estimated Risk Score: {risk:.1f} (placeholder ML logic)"


# ---------------------------------------------------------
# Gradio Interfaces
# ---------------------------------------------------------

# Imaging Tab
image_interface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(label="Upload Chest X-ray"),
    outputs=[
        gr.Textbox(label="Status"),
        gr.Image(label="AI Heatmap Output")
    ],
    title="ContinuumDx — Imaging Module",
    description="Upload a chest X-ray to view placeholder AI inference."
)

# Cough Audio Tab
audio_interface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(
        sources=["upload"],
        type="filepath",
        label="Upload Cough / Respiratory Audio"
    ),
    outputs=gr.Textbox(label="Audio Analysis Result"),
    title="Cough Sound Analysis",
    description="Upload cough audio to simulate AI cough analysis."
)

# Symptoms Tab
symptom_interface = gr.Interface(
    fn=process_tabular,
    inputs=[
        gr.CheckboxGroup(
            ["Persistent Cough", "Breathlessness", "Fever", "Fatigue"],
            label="Symptoms"
        ),
        gr.Number(label="Age"),
        gr.Checkbox(label="Smoker?")
    ],
    outputs=gr.Textbox(label="Risk Score"),
    title="Symptoms + Risk Assessment",
    description="Enter symptoms and risk factors for structured health screening."
)

# Combine Tabs
app = gr.TabbedInterface(
    [image_interface, audio_interface, symptom_interface],
    ["Imaging", "Cough Audio", "Symptoms"]
)

# Run App
if __name__ == "__main__":
    app.launch(share=True)