report_generator.py

"""
Report Generator for NeuroSAM3.

Generates structured clinical and research reports from segmentation findings
using LLM reasoning (via HF Inference API).
"""

from typing import Optional, Dict, Any
from logger_config import logger
from config import DEFAULT_REPORT_STYLE

CLINICAL_REPORT_TEMPLATE = """Based on the following neuroimaging findings, generate a structured radiology report.

**Findings from AI analysis:**
{findings}

**Clinical Context:**
{clinical_context}

**Report Style:** {report_style}

Generate a structured report with these sections:
1. TECHNIQUE: Brief description of imaging modality and analysis method
2. FINDINGS: Detailed description of segmentation results and measurements
3. IMPRESSION: Summary of key findings (2-3 bullet points)
4. MEASUREMENTS: Quantitative data (area, volume, intensity if available)
5. LIMITATIONS: Note this is AI-assisted and requires clinical correlation

IMPORTANT: Frame all findings as "AI-assisted observations" requiring clinical correlation.
Never state definitive diagnoses."""

RESEARCH_REPORT_TEMPLATE = """Based on the following neuroimaging analysis results, generate a research summary.

**Analysis Results:**
{findings}

**Study Context:**
{clinical_context}

Generate a structured research report with:
1. METHODS: Analysis pipeline and models used
2. RESULTS: Quantitative findings with statistics
3. OBSERVATIONS: Qualitative findings
4. DATA QUALITY: Confidence scores and limitations
5. SUGGESTED NEXT STEPS: Further analyses recommended"""


def generate_report(
    findings: str,
    style: str = DEFAULT_REPORT_STYLE,
    clinical_context: str = "",
    api_client=None,
) -> str:
    """
    Generate a structured report from findings.

    Args:
        findings: Summary of segmentation/analysis findings
        style: "radiology" | "neurosurgery" | "research"
        clinical_context: Optional patient history or study context
        api_client: HFInferenceAPI instance (optional, uses global if None)

    Returns:
        Formatted report string
    """
    if not findings:
        return "No findings to report. Please run segmentation first."

    if not clinical_context:
        clinical_context = "Not provided"

    # Select template
    if style == "research":
        template = RESEARCH_REPORT_TEMPLATE
    else:
        template = CLINICAL_REPORT_TEMPLATE

    prompt = template.format(
        findings=findings,
        clinical_context=clinical_context,
        report_style=style,
    )

    # Try LLM generation
    if api_client and api_client.is_available:
        try:
            response = api_client.chat(
                messages=[{"role": "user", "content": prompt}],
                system_prompt="You are a neuroradiology reporting assistant. Generate structured, professional reports.",
                max_tokens=1500,
                temperature=0.2,
            )
            if response:
                return response
        except Exception as e:
            logger.warning(f"LLM report generation failed: {e}")

    # Fallback: template-based report without LLM
    return _generate_template_report(findings, style, clinical_context)


def _generate_template_report(
    findings: str,
    style: str,
    clinical_context: str,
) -> str:
    """Generate a basic template report without LLM (fallback)."""
    if style == "research":
        return f"""## Research Analysis Report

### Methods
- Segmentation: SAM3 / MedSAM (text-prompted / bounding-box)
- Classification: BiomedCLIP zero-shot
- Platform: NeuroSAM3 (HuggingFace Spaces)

### Results
{findings}

### Study Context
{clinical_context}

### Data Quality
- AI-assisted analysis — results should be validated
- Confidence scores reported per segmentation

### Suggested Next Steps
- Cross-validate with manual annotations
- Compare across subjects for statistical significance
- Export NIFTI masks for volumetric analysis in external tools

---
*Generated by NeuroSAM3 | AI-assisted — requires expert validation*"""

    else:  # radiology / neurosurgery
        return f"""## AI-Assisted Neuroimaging Report

### Technique
Automated segmentation and analysis using NeuroSAM3 platform.
Models: SAM3 (general), MedSAM (pathology), BiomedCLIP (classification).

### Findings
{findings}

### Clinical Context
{clinical_context}

### Impression
- AI-assisted analysis completed
- Findings require clinical correlation
- See measurements above for quantitative data

### Limitations
- This is an AI-assisted report and does NOT constitute a medical diagnosis
- All findings require correlation with clinical presentation
- Segmentation accuracy varies by structure and image quality

---
*Generated by NeuroSAM3 | NOT a diagnostic report — requires clinical correlation*"""


def generate_comparison_report(
    results: list,
    models_used: list,
) -> str:
    """Generate a model comparison report."""
    report = "## Model Comparison Report\n\n"
    report += "| Model | Detected | Area (px) | Score |\n"
    report += "|-------|----------|-----------|-------|\n"

    for result, model in zip(results, models_used):
        if result and "mask" in result:
            import numpy as np
            area = int(np.sum(result["mask"]))
            score = result.get("score", "N/A")
            report += f"| {model} | Yes | {area:,} | {score:.3f} |\n"
        else:
            report += f"| {model} | No | 0 | - |\n"

    report += "\n*Lower threshold models may detect more but with lower specificity.*\n"
    return report