"""
TB-Guard-XAI · End-to-end pipeline example
Image → Preprocess → Ensemble → Grad-CAM++ → Mistral clinical report
"""

import cv2
import torch
import numpy as np

from preprocessing import LungPreprocessor, get_val_transforms
from ensemble_models import load_ensemble
from gradcam import GradCAMPlusPlus
from mistral_explainer import MistralExplainer


def run(image_path: str, checkpoint: str = None, api_key: str = None):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # 1. Load model
    model = load_ensemble(checkpoint, device=device)

    # 2. Preprocess X-ray
    preprocessor = LungPreprocessor(image_size=224)
    image_np = preprocessor.preprocess(image_path, segment_lung=True)   # (224,224) uint8

    transforms = get_val_transforms(224)
    tensor = transforms(image=image_np)["image"].unsqueeze(0).to(device)  # (1,1,224,224)

    # 3. Grad-CAM++ → structured output
    gcam   = GradCAMPlusPlus(model, device=device)
    output = gcam.explain(
        tensor,
        threshold=0.5,
        n_mc_samples=20,
        # Wire in ExplainabilityValidator results if available:
        radiologist_agreement=0.75,
        n_radiologists=4,
        explanation_valid_votes=3,
    )

    # 4. Save overlay image
    overlay = gcam.heatmap_overlay(image_np, np.zeros((224,224), dtype=np.float32))
    cv2.imwrite("gradcam_overlay.png", overlay)
    print(f"Overlay saved → gradcam_overlay.png")

    # 5. Mistral clinical report
    explainer = MistralExplainer(api_key=api_key)
    report    = explainer.generate_report(output)
    report.print_report()
    report.save("tb_report.json")

    return report


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("image",      help="Path to chest X-ray (PNG/JPG/DICOM)")
    p.add_argument("--checkpoint", default=None, help="Path to ensemble .pth weights")
    p.add_argument("--api-key",    default=None, help="Mistral API key (or set MISTRAL_API_KEY)")
    args = p.parse_args()
    run(args.image, args.checkpoint, args.api_key)