src/main.py

import os

import gradio as gr
from models.inference import (
    run_inference,
    get_available_model_choices,
    get_default_model_key,
    get_model_label,
)
from models.metadata_builder import build_metadata_csv

GROUP_CHOICES = [
    ("Demographics", "demographics"),
    ("Clinical History", "history"),
    ("Symptoms", "symptoms"),
    ("Lesion Geometry", "lesion_geometry"),
]
DEFAULT_GROUPS = ["demographics", "symptoms", "lesion_geometry"]
REGION_CHOICES = ["HEAD", "NECK", "BACK", "ARM", "LEG", "TORSO"]
GENDER_CHOICES = ["MALE", "FEMALE"]

MODEL_CHOICES = get_available_model_choices()
DEFAULT_MODEL_KEY = get_default_model_key()

PAPER_URL = os.environ.get("PAPER_URL", "").strip()
PAPER_TITLE = os.environ.get(
    "PAPER_TITLE",
    "RG-DermNet: Multimodal Skin Lesion Explainability"
).strip()
PAPER_DESCRIPTION = os.environ.get(
    "PAPER_DESCRIPTION",
    "This Space accompanies the proposed multimodal framework and allows real-time "
    "inspection of how clinical metadata affects prediction behavior and GradCAM++ attention maps."
).strip()

PAPER_FILE = os.environ.get("PAPER_FILE", "paper.pdf").strip()
PAPER_EXISTS = os.path.exists(PAPER_FILE)

custom_css = """
.gradio-container {
    background: linear-gradient(180deg, #08111f 0%, #0b0f19 100%) !important;
    color: #e8eef8 !important;
}

.main-shell {
    max-width: 1280px;
    margin: 0 auto;
}

.hero {
    padding: 28px 24px 18px 24px;
    border: 1px solid #24364d;
    border-radius: 18px;
    background: linear-gradient(135deg, rgba(17, 28, 46, 0.95), rgba(8, 13, 24, 0.95));
    margin-bottom: 18px;
}

.hero h1 {
    margin: 0 0 10px 0;
    font-size: 2.1rem;
    line-height: 1.2;
}

.hero p {
    margin: 0;
    color: #d2dceb;
    line-height: 1.65;
    font-size: 1rem;
}

.badge-row {
    display: flex;
    flex-wrap: wrap;
    gap: 8px;
    margin-top: 14px;
}

.badge {
    padding: 6px 10px;
    border-radius: 999px;
    background: #12253f;
    border: 1px solid #2f527c;
    color: #beddff;
    font-size: 0.9rem;
}

.section-card {
    border: 1px solid #24364d;
    border-radius: 18px;
    background: rgba(11, 18, 31, 0.92);
    padding: 18px 20px;
    margin-bottom: 18px;
}

.section-card h2,
.section-card h3 {
    margin-top: 0;
}

.paper-card {
    padding: 16px 18px;
    border: 1px solid #2d4f7c;
    border-radius: 14px;
    background: linear-gradient(135deg, rgba(18, 34, 58, 0.95), rgba(10, 16, 28, 0.95));
}

.paper-card h3 {
    margin: 0 0 8px 0;
    font-size: 1.05rem;
}

.paper-card p {
    margin: 0 0 8px 0;
    color: #d3deee;
    line-height: 1.55;
}

.paper-card a {
    color: #8bc4ff;
    text-decoration: none;
    font-weight: 600;
}

.paper-card a:hover {
    text-decoration: underline;
}

.pipeline-box {
    padding: 14px;
    border-radius: 14px;
    border: 1px solid #2b405e;
    background: #0d1727;
    text-align: center;
    min-height: 120px;
    display: flex;
    flex-direction: column;
    justify-content: center;
}

.pipeline-box h3 {
    margin-bottom: 8px;
}

.pipeline-box p {
    margin: 0;
    color: #ced9ea;
    line-height: 1.5;
}

.pipeline-arrow {
    text-align: center;
    font-size: 1.6rem;
    color: #8bc4ff;
    padding-top: 38px;
    font-weight: 700;
}

.demo-panel {
    border: 1px solid #2d3748;
    padding: 16px;
    border-radius: 16px;
    background: #111a29;
}

.predict-btn {
    background: #3182ce !important;
    color: white !important;
    font-weight: bold !important;
    border: none !important;
}

.predict-btn:hover {
    background: #4299e1 !important;
}

.soft-text {
    color: #b9c7da;
    line-height: 1.6;
}

.footer-note {
    font-size: 0.95rem;
    color: #b7c6db;
    line-height: 1.6;
}

.paper-frame-wrap {
    border: 1px solid #2d4f7c;
    border-radius: 14px;
    overflow: hidden;
    background: #0b1321;
}

.paper-frame {
    width: 100%;
    height: 900px;
    border: none;
    background: white;
}

.muted-divider {
    opacity: 0.25;
    margin: 10px 0 14px 0;
}
"""


def build_paper_card():
    link_html = ""
    if PAPER_URL:
        link_html += (
            f'<p><a href="{PAPER_URL}" target="_blank" rel="noopener noreferrer">'
            f'Open paper link</a></p>'
        )

    if PAPER_EXISTS:
        link_html += '<p><a href="/file=paper.pdf" target="_blank" rel="noopener noreferrer">Open embedded PDF in new tab</a></p>'

    if not link_html:
        link_html = "<p>No external paper link configured yet.</p>"

    return f"""
    <div class="paper-card">
        <h3>📄 Associated Paper</h3>
        <p><strong>{PAPER_TITLE}</strong></p>
        <p>{PAPER_DESCRIPTION}</p>
        {link_html}
    </div>
    """


def build_hero():
    return """
    <div class="hero">
        <h1>🔬  RG-DermNet: A Multimodal Attention-Based Model with Residual Block Usage for Skin Lesion Classification </h1>
        <p>
            This interactive scientific demo presents a multimodal skin lesion analysis system that combines
            clinical images and patient metadata to generate predictions and GradCAM++ explanations.
            The interface allows real-time inspection of how metadata influences model attention and
            diagnostic behavior.
        </p>
        <div class="badge-row">
            <span class="badge">Clinical Image + Metadata</span>
            <span class="badge">Multimodal Attention</span>
            <span class="badge">GradCAM++ Explainability</span>
            <span class="badge">Interactive Paper Demo</span>
        </div>
    </div>
    """


def build_paper_embed():
    if not PAPER_EXISTS:
        return """
        <div class="paper-card">
            <h3>Paper preview unavailable</h3>
            <p>
                The file <strong>paper.pdf</strong> was not found in the repository root.
                Add it to enable in-Space preview.
            </p>
        </div>
        """

    return """
    <div class="paper-frame-wrap">
        <iframe src="/file=paper.pdf" class="paper-frame"></iframe>
    </div>
    """


def format_groups(enabled_groups):
    if not enabled_groups:
        return "No metadata group selected."
    label_map = dict(GROUP_CHOICES)
    return " | ".join([label_map.get(g, g) for g in enabled_groups])


def safe_bool(value):
    return bool(value)


def build_values_dict(age, gender, region, diameter1, diameter2, itch, grew, hurt, changed, bleed, elevation):
    return {
        "age": age,
        "gender": gender,
        "region": region,
        "diameter_1": diameter1,
        "diameter_2": diameter2,
        "itch": safe_bool(itch),
        "grew": safe_bool(grew),
        "hurt": safe_bool(hurt),
        "changed": safe_bool(changed),
        "bleed": safe_bool(bleed),
        "elevation": safe_bool(elevation),
    }


def build_metadata_preview(enabled_groups, age, gender, region, diameter1, diameter2, itch, grew, hurt, changed, bleed, elevation):
    values = build_values_dict(age, gender, region, diameter1, diameter2, itch, grew, hurt, changed, bleed, elevation)
    metadata_csv = build_metadata_csv(values, enabled_groups)
    groups_text = format_groups(enabled_groups)
    return metadata_csv, groups_text


def validate_inputs(image, enabled_groups, age, diameter1, diameter2):
    if image is None:
        raise gr.Error("Please upload a dermoscopic image first.")
    if not enabled_groups:
        raise gr.Error("Please select at least one metadata group.")
    if age is None or age < 0:
        raise gr.Error("Age must be a valid non-negative number.")
    if diameter1 is None or diameter1 < 0:
        raise gr.Error("Diameter 1 must be a valid non-negative number.")
    if diameter2 is None or diameter2 < 0:
        raise gr.Error("Diameter 2 must be a valid non-negative number.")


def gradio_predict(image, selected_model_key, enabled_groups, age, gender, region, diameter1, diameter2, itch, grew, hurt, changed, bleed, elevation):
    validate_inputs(image, enabled_groups, age, diameter1, diameter2)
    values = build_values_dict(age, gender, region, diameter1, diameter2, itch, grew, hurt, changed, bleed, elevation)
    metadata_csv = build_metadata_csv(values, enabled_groups)

    try:
        heatmap_img, prediction_text = run_inference(image, metadata_csv, selected_model_key)
    except RuntimeError as exc:
        raise gr.Error(str(exc)) from exc

    groups_text = format_groups(enabled_groups)
    model_text = get_model_label(selected_model_key)
    pretty_prediction = (
        f"### 🩺 Prediction Result\n\n"
        f"**Selected model:** {model_text}\n\n"
        f"**Active groups:** {groups_text}\n\n"
        f"**Model output:**\n{prediction_text}"
    )
    return image, heatmap_img, pretty_prediction, metadata_csv, groups_text


def clear_all():
    default_model = DEFAULT_MODEL_KEY
    if default_model is None and MODEL_CHOICES:
        default_model = MODEL_CHOICES[0][1]

    return (
        None, DEFAULT_GROUPS, default_model, 55.0, "FEMALE", "NECK", 6.0, 5.0,
        False, False, False, False, False, False,
        None, None, "### Prediction Result\n\nRun the model to see the output here.",
        "", format_groups(DEFAULT_GROUPS)
    )


with gr.Blocks(
    title="Skin Lesion Explainability",
    theme=gr.themes.Default(primary_hue="blue"),
    css=custom_css,
) as demo:
    with gr.Column(elem_classes="main-shell"):
        gr.HTML(build_hero())

        with gr.Group(elem_classes="section-card"):
            gr.Markdown("## 🧪 Interactive Demonstration")
            gr.Markdown(
                "Modify the metadata, choose a model variant, and inspect how the attention map changes.",
                elem_classes="soft-text",
            )

            with gr.Row():
                with gr.Column(scale=1, elem_classes="demo-panel"):
                    gr.Markdown("### 📥 Input Data")

                    image_input = gr.Image(type="pil", label="Dermoscopic Image", height=320)

                    group_selector = gr.CheckboxGroup(
                        choices=GROUP_CHOICES,
                        value=DEFAULT_GROUPS,
                        label="Enable Metadata Groups"
                    )

                    model_selector = gr.Dropdown(
                        choices=MODEL_CHOICES,
                        value=DEFAULT_MODEL_KEY if DEFAULT_MODEL_KEY is not None else None,
                        label="Attention Mechanism Model",
                        info="Choose which pretrained attention mechanism/model to run.",
                    )

                    with gr.Accordion("👤 Demographics", open=True):
                        age = gr.Number(label="Age", value=55, precision=0)
                        with gr.Row():
                            gender = gr.Dropdown(GENDER_CHOICES, value="FEMALE", label="Gender")
                            region = gr.Dropdown(REGION_CHOICES, value="NECK", label="Region")

                    with gr.Accordion("📏 Lesion Geometry", open=False):
                        with gr.Row():
                            diameter1 = gr.Number(label="Diameter 1", value=6)
                            diameter2 = gr.Number(label="Diameter 2", value=5)

                    with gr.Accordion("🚩 Symptoms", open=False):
                        with gr.Row():
                            itch = gr.Checkbox(label="Itch")
                            grew = gr.Checkbox(label="Grew")
                            hurt = gr.Checkbox(label="Hurt")
                        with gr.Row():
                            changed = gr.Checkbox(label="Changed")
                            bleed = gr.Checkbox(label="Bleed")
                            elevation = gr.Checkbox(label="Elevation")

                    with gr.Row():
                        clear_btn = gr.Button("Clear", variant="secondary")
                        run_btn = gr.Button("Generate GradCAM++", variant="primary", elem_classes="predict-btn")

                with gr.Column(scale=2, elem_classes="demo-panel"):
                    gr.Markdown("### 📊 Analysis Dashboard")

                    with gr.Row():
                        original_img_out = gr.Image(label="Original Lesion", interactive=False)
                        heatmap_out = gr.Image(label="Attention Map (GradCAM++)", interactive=False)

                    with gr.Group():
                        prediction_out = gr.Markdown(
                            value="### Prediction Result\n\nRun the model to see the output here."
                        )

                    with gr.Accordion("📋 System Metadata Details", open=False):
                        active_groups_text = gr.Textbox(label="Active Groups", interactive=False)
                        metadata_preview = gr.Textbox(label="Final CSV Input", lines=6, interactive=False)

        with gr.Group(elem_classes="section-card"):
            gr.Markdown("## 📚 Notes for Readers")
            gr.Markdown(
                """
- This demo is intended as a qualitative companion to the paper.
- Users can inspect how metadata groups influence model behavior and attention maps.
- The available models correspond to pretrained multimodal attention-based variants.
- For best scientific use, this Space should be interpreted together with the associated manuscript.
""",
                elem_classes="footer-note",
            )

            gr.Markdown("## 📎 Citation and Reproducibility")
            gr.Markdown(
                """
If you reference this demo in a paper, thesis, or presentation, cite the associated manuscript
and include the Hugging Face Space as supplementary interactive material.

BibTeX: 

@inproceedings{rocha2026rgdermnet,
  title     = {RG-DermNet: A Multimodal Attention-Based Model with Residual Block Usage for Skin Lesion Classification},
  author    = {Rocha, Wyctor F. and Bouzon, Pedro H. G. and Ramos, Lucas A. and Pacheco, Andre G. C. and Souza Jr., Luis A.},
  booktitle = {International Joint Conference on Neural Networks (IJCNN)},
  year      = {2026},
  note      = {Accepted}
}
""",
                elem_classes="footer-note",
            )

    preview_inputs = [
        group_selector, age, gender, region, diameter1, diameter2,
        itch, grew, hurt, changed, bleed, elevation
    ]

    for component in preview_inputs:
        component.change(
            fn=build_metadata_preview,
            inputs=preview_inputs,
            outputs=[metadata_preview, active_groups_text]
        )

    run_btn.click(
        fn=gradio_predict,
        inputs=[image_input, model_selector] + preview_inputs,
        outputs=[original_img_out, heatmap_out, prediction_out, metadata_preview, active_groups_text]
    )

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[
            image_input, group_selector, model_selector, age, gender, region, diameter1, diameter2,
            itch, grew, hurt, changed, bleed, elevation,
            original_img_out, heatmap_out, prediction_out, metadata_preview, active_groups_text
        ]
    )