explainability_validator.py

"""
Explainability Validation Module
Issue #28: Validate that Grad-CAM explanations align with actual pathology
"""

import logging
import json
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
from datetime import datetime
import cv2

logger = logging.getLogger("tb_guard_explainability")


@dataclass
class ExplainabilityCase:
    """Single explainability validation case"""
    case_id: str
    xray_path: str
    prediction: str  # "TB" or "Normal"
    ground_truth: str  # "TB" or "Normal"
    prediction_correct: bool
    
    # Grad-CAM metrics
    gradcam_region: str  # "upper", "middle", "lower", "diffuse"
    gradcam_intensity: float  # 0-1, peak intensity
    
    # Radiologist validation
    radiologist_votes: Dict[str, bool]  # radiologist_id -> "does heatmap show pathology?"
    radiologist_agreement: float  # 0-1, fraction in agreement


class ExplainabilityValidator:
    """
    Validates that model explanations (Grad-CAM) are trustworthy
    Issue #28: Ensure radiologists trust the model's visual explanations
    """
    
    def __init__(self, output_dir: str = "explainability_validation"):
        self.output_dir = Path(output_dir)
        try:
            self.output_dir.mkdir(exist_ok=True)
        except PermissionError:
            import tempfile
            self.output_dir = Path(tempfile.gettempdir()) / "tb_guard_explainability"
            self.output_dir.mkdir(exist_ok=True)
        self.cases: List[ExplainabilityCase] = []
        
    def create_validation_case(
        self,
        case_id: str,
        xray_path: str,
        model_prediction: str,
        ground_truth: str,
        gradcam_heatmap: np.ndarray
    ) -> Dict:
        """
        Create a case for radiologist validation
        
        Args:
            case_id: Unique case identifier
            xray_path: Path to original X-ray image
            model_prediction: Model's prediction ("TB" or "Normal")
            ground_truth: Actual diagnosis
            gradcam_heatmap: Grad-CAM heatmap (H x W array, 0-1)
        
        Returns:
            Validation form with image, heatmap, and prediction
        """
        # Verify heatmap
        if gradcam_heatmap is None:
            logger.warning(f"No heatmap for case {case_id}")
            return {"error": "No heatmap available"}
        
        # Save heatmap for radiologist review
        heatmap_colored = cv2.applyColorMap(
            (gradcam_heatmap * 255).astype(np.uint8),
            cv2.COLORMAP_JET
        )
        heatmap_path = self.output_dir / f"{case_id}_heatmap.png"
        cv2.imwrite(str(heatmap_path), heatmap_colored)
        
        # Calculate heatmap statistics
        peak_intensity = float(np.max(gradcam_heatmap))
        h, w = gradcam_heatmap.shape
        upper_intensity = np.mean(gradcam_heatmap[:h//3])
        middle_intensity = np.mean(gradcam_heatmap[h//3:2*h//3])
        lower_intensity = np.mean(gradcam_heatmap[2*h//3:])
        
        # Determine dominant region (before asking radiologists)
        intensities = {"upper": upper_intensity, "middle": middle_intensity, "lower": lower_intensity}
        dominant_region = max(intensities, key=intensities.get)
        
        return {
            "case_id": case_id,
            "xray_path": xray_path,
            "heatmap_path": str(heatmap_path),
            "model_prediction": model_prediction,
            "ground_truth": ground_truth,
            "prediction_correct": model_prediction == ground_truth,
            "heatmap_stats": {
                "peak_intensity": peak_intensity,
                "dominant_region": dominant_region,
                "upper_intensity": upper_intensity,
                "middle_intensity": middle_intensity,
                "lower_intensity": lower_intensity
            },
            "validation_question": (
                "Does this heatmap highlight the area you would focus on when diagnosing TB? "
                "Answer YES if the highlighted region(s) contain actual pathology. "
                "Answer NO if the highlight is on artifacts, normal anatomy, or wrong regions."
            )
        }
    
    def record_radiologist_feedback(
        self,
        case_id: str,
        radiologist_id: str,
        is_valid: bool,
        confidence: float,
        notes: str = ""
    ):
        """
        Record a radiologist's assessment of the Grad-CAM explanation
        
        Args:
            case_id: Case identifier
            radiologist_id: Radiologist identifier
            is_valid: True if radiologist agrees heatmap is correct
            confidence: 0-1, radiologist confidence in their judgment
            notes: Optional notes
        """
        feedback_file = self.output_dir / f"{case_id}_feedback.jsonl"
        feedback = {
            "timestamp": datetime.utcnow().isoformat(),
            "radiologist_id": radiologist_id,
            "is_valid": is_valid,
            "confidence": confidence,
            "notes": notes
        }
        with open(feedback_file, "a") as f:
            f.write(json.dumps(feedback) + "\n")
    
    def finalize_case(
        self,
        case_id: str,
        prediction: str,
        ground_truth: str,
        gradcam_region: str,
        gradcam_intensity: float
    ) -> ExplainabilityCase:
        """
        Finalize a case after radiologist feedback is collected
        
        Returns:
            ExplainabilityCase with aggregated metrics
        """
        # Read radiologist votes
        feedback_file = self.output_dir / f"{case_id}_feedback.jsonl"
        votes = {}
        all_valid = []
        
        if feedback_file.exists():
            with open(feedback_file, "r") as f:
                for line in f:
                    feedback = json.loads(line)
                    radiologist_id = feedback["radiologist_id"]
                    votes[radiologist_id] = feedback["is_valid"]
                    all_valid.append(feedback["is_valid"])
        
        # Calculate agreement (fraction of radiologists who thought explanation was valid)
        agreement = np.mean(all_valid) if all_valid else 0.0
        
        case = ExplainabilityCase(
            case_id=case_id,
            xray_path=f"case_{case_id}.png",
            prediction=prediction,
            ground_truth=ground_truth,
            prediction_correct=prediction == ground_truth,
            gradcam_region=gradcam_region,
            gradcam_intensity=gradcam_intensity,
            radiologist_votes=votes,
            radiologist_agreement=agreement
        )
        
        self.cases.append(case)
        return case
    
    def generate_report(self) -> Dict:
        """
        Generate explainability validation report
        Issue #28: Complete metrics on whether explanations are trustworthy
        """
        if not self.cases:
            return {"status": "No validation cases"}
        
        # Overall statistics
        total_cases = len(self.cases)
        correct_predictions = sum(1 for c in self.cases if c.prediction_correct)
        valid_explanations = sum(1 for c in self.cases if c.radiologist_agreement > 0.67)  # 2/3 agreement
        
        # Accuracy
        prediction_accuracy = correct_predictions / total_cases if total_cases > 0 else 0
        
        # Explanation validity (radiologists agree with Grad-CAM)
        explanation_validity = valid_explanations / total_cases if total_cases > 0 else 0
        
        # Cross-tabulation: (model_correct, explanation_valid)
        model_correct_expl_valid = sum(
            1 for c in self.cases
            if c.prediction_correct and c.radiologist_agreement > 0.67
        )
        model_correct_expl_invalid = sum(
            1 for c in self.cases
            if c.prediction_correct and c.radiologist_agreement <= 0.67
        )
        model_wrong_expl_valid = sum(
            1 for c in self.cases
            if not c.prediction_correct and c.radiologist_agreement > 0.67
        )
        model_wrong_expl_invalid = sum(
            1 for c in self.cases
            if not c.prediction_correct and c.radiologist_agreement <= 0.67
        )
        
        # Per-region analysis
        region_performance = {}
        for region in ["upper", "middle", "lower", "diffuse"]:
            cases_in_region = [c for c in self.cases if c.gradcam_region == region]
            if cases_in_region:
                region_valid = sum(1 for c in cases_in_region if c.radiologist_agreement > 0.67)
                region_performance[region] = {
                    "count": len(cases_in_region),
                    "explanation_validity": region_valid / len(cases_in_region)
                }
        
        # Issues with explanations
        issues = []
        
        # Issue: Model correct but explanation invalid
        if model_correct_expl_invalid > 0:
            issues.append({
                "type": "CORRECT_PREDICTION_INVALID_EXPLANATION",
                "count": model_correct_expl_invalid,
                "severity": "MEDIUM",
                "description": "Model made correct prediction but radiologists don't trust the explanation. "
                               "Indicates Grad-CAM is not capturing clinically relevant features."
            })
        
        # Issue: Model wrong but explanation was "valid"
        if model_wrong_expl_valid > 0:
            issues.append({
                "type": "INCORRECT_PREDICTION_VALID_EXPLANATION",
                "count": model_wrong_expl_valid,
                "severity": "HIGH",
                "description": "Model made wrong prediction but Grad-CAM highlighted plausible areas. "
                               "Indicates explanations can be misleading even when wrong."
            })
        
        # Issue: Overall low explanation validity
        if explanation_validity < 0.75:
            issues.append({
                "type": "LOW_OVERALL_EXPLANATION_VALIDITY",
                "validity_percent": explanation_validity * 100,
                "severity": "HIGH",
                "description": "Less than 75% of explanations are trusted by radiologists. "
                               "Model should not be used in production until addressed."
            })
        
        report = {
            "timestamp": datetime.utcnow().isoformat(),
            "total_cases_validated": total_cases,
            "metrics": {
                "prediction_accuracy": prediction_accuracy,
                "explanation_validity": explanation_validity,
                "explanation_validity_percent": explanation_validity * 100
            },
            "confusion_matrix": {
                "model_correct_expl_valid": model_correct_expl_valid,
                "model_correct_expl_invalid": model_correct_expl_invalid,
                "model_wrong_expl_valid": model_wrong_expl_valid,
                "model_wrong_expl_invalid": model_wrong_expl_invalid
            },
            "by_region": region_performance,
            "issues": issues,
            "recommendation": (
                "✓ APPROVED FOR PRODUCTION" if explanation_validity >= 0.80 and len(issues) == 0
                else "⚠ CONDITIONAL APPROVAL - Address issues before deployment"
                if explanation_validity >= 0.75
                else "✗ NOT APPROVED - Explanation validity too low"
            )
        }
        
        # Save report
        report_path = self.output_dir / "explainability_report.json"
        with open(report_path, "w") as f:
            json.dump(report, f, indent=2)
        
        logger.info(f"Explainability report saved to {report_path}")
        return report


# Example usage in evaluation pipeline:
# validator = ExplainabilityValidator()
#
# for test_case in test_set:
#     prediction = model.predict(test_case.image)
#     gradcam_heatmap = generate_gradcam(test_case.image)
#     
#     # Prepare validation form for radiologists
#     form = validator.create_validation_case(
#         case_id=test_case.id,
#         xray_path=test_case.path,
#         model_prediction=prediction,
#         ground_truth=test_case.label,
#         gradcam_heatmap=gradcam_heatmap
#     )
#     # Send form to radiologist panel
#
# # After radiologists provide feedback:
# radiologist_feedback = get_feedback_from_panel()
# for feedback in radiologist_feedback:
#     validator.record_radiologist_feedback(
#         case_id=feedback.case_id,
#         radiologist_id=feedback.radiologist_id,
#         is_valid=feedback.is_valid,
#         confidence=feedback.confidence
#     )
#
# # Generate final report
# report = validator.generate_report()
# print(report["recommendation"])