medical_templates.py

DISEASE_TEMPLATES = {
    "anemia": {
        "disease_name": "Iron-Deficiency Anemia",
        "visual_signs": {
            "conjunctival_pallor": 0.85,        # Pale inner eyelid
            "oral_mucosal_pallor": 0.80,        # Pale mouth
            "glossitis": 0.70,                  # Smooth tongue
            "angular_cheilitis": 0.65,          # Cracked corners
            "facial_pallor": 0.60,
        },
        "clinical_rules": [
            {"condition": "conjunctival_pallor > 0.75 AND oral_mucosal_pallor > 0.70", "boost": 0.15},
            {"condition": "glossitis > 0.60 AND conjunctival_pallor > 0.70", "boost": 0.20},
            {"condition": "age >= 40", "boost": 0.05},
        ],
        "explanation_en": "Your eyes and mouth show signs of paleness, suggesting iron-deficiency anemia. Get hemoglobin blood test.",
        "explanation_hi": "आपकी आँखों और मुँह में पीलापन के संकेत दिख रहे हैं। हीमोग्लोबिन परीक्षण करवाएं।",
        "urgency": "routine",
        "follow_up_tests": ["Hemoglobin (Hb)", "Serum Iron", "Ferritin", "TIBC"],
        "recommendations": ["Get blood tests", "Increase iron-rich foods", "Consult doctor"]
    },

    "jaundice": {
        "disease_name": "Jaundice / Liver Disease",
        "visual_signs": {
            "scleral_icterus": 0.95,            # Yellow sclera (strongest indicator)
            "palmar_erythema": 0.80,            # Red palms
            "xanthelasma": 0.70,                # Yellow eyelid plaques
            "facial_jaundice": 0.90,            # Yellow face
        },
        "clinical_rules": [
            {"condition": "scleral_icterus > 0.85", "boost": 0.25, "urgency_upgrade": "urgent"},
            {"condition": "scleral_icterus > 0.75 AND palmar_erythema > 0.70", "boost": 0.30, "urgency": "urgent"},
        ],
        "explanation_en": "Yellow eyes suggest liver disease. URGENT: Get liver function tests immediately.",
        "explanation_hi": "पीली आँखें लीवर की बीमारी का संकेत देती हैं। तुरंत लीवर फंक्शन टेस्ट करवाएं।",
        "urgency": "urgent",
        "follow_up_tests": ["Bilirubin (Total & Direct)", "ALT", "AST", "ALP", "Albumin", "Hepatitis Panel"],
        "recommendations": ["URGENT: Seek medical attention", "Get liver function tests", "Avoid alcohol"]
    },

    "stroke_risk": {
        "disease_name": "Stroke Risk / Neurological Signs",
        "visual_signs": {
            "facial_asymmetry": 0.90,           # One-sided drooping
            "mouth_drooping": 0.85,             # Mouth not aligned
            "eye_deviation": 0.80,              # Eyes not aligned
        },
        "clinical_rules": [
            {"condition": "facial_asymmetry > 0.80", "boost": 0.25, "urgency_upgrade": "urgent"},
            {"condition": "facial_asymmetry > 0.75 AND age >= 50", "boost": 0.25, "urgency_upgrade": "urgent"},
        ],
        "explanation_en": "EMERGENCY: Face drooping may indicate STROKE. Call ambulance immediately!",
        "explanation_hi": "आपातकाल: चेहरे में लकवा स्ट्रोक का संकेत दे सकता है। तुरंत एंबुलेंस बुलाएं!",
        "urgency": "urgent",
        "follow_up_tests": ["CT/MRI Brain", "ECG", "Blood glucose", "Troponin"],
        "recommendations": ["CALL AMBULANCE IMMEDIATELY", "Note time symptoms started", "Do NOT eat/drink"]
    },

    "hypoxia": {
        "disease_name": "Hypoxia / Respiratory Disease (COPD/Asthma)",
        "visual_signs": {
            "cyanosis_lips": 0.90,              # Bluish lips
            "cyanosis_nails": 0.85,             # Bluish nails
            "nail_clubbing": 0.70,              # Clubbed nails (chronic disease sign)
            "facial_cyanosis": 0.80,
        },
        "clinical_rules": [
            {"condition": "cyanosis_lips > 0.80", "boost": 0.25, "urgency_upgrade": "moderate"},
            {"condition": "age >= 40 AND nail_clubbing > 0.60", "boost": 0.15},
        ],
        "explanation_en": "Blue lips indicate low blood oxygen. Get SpO2 test immediately. If < 90%, go to hospital.",
        "explanation_hi": "नीले होंठ कम ऑक्सीजन का संकेत देते हैं। तुरंत SpO2 परीक्षण करवाएं।",
        "urgency": "moderate",
        "follow_up_tests": ["SpO2 (Pulse Oximetry)", "PFT (FEV1, FVC)", "Chest X-ray", "ABG"],
        "recommendations": ["Get oxygen saturation test", "If SpO2 < 90%, go to hospital immediately", "Quit smoking"]
    },

    "conjunctivitis": {
        "disease_name": "Conjunctivitis / Eye Infection",
        "visual_signs": {
            "conjunctival_redness": 0.90,       # Red conjunctiva
            "eye_discharge": 0.80,              # Purulent/watery discharge
            "conjunctival_swelling": 0.75,
        },
        "clinical_rules": [
            {"condition": "conjunctival_redness > 0.80", "boost": 0.20},
            {"condition": "conjunctival_redness > 0.75 AND eye_discharge > 0.70", "boost": 0.25},
        ],
        "explanation_en": "Red eyes suggest infection. Wash with clean water, use prescribed drops. See doctor if persists >7 days.",
        "explanation_hi": "लाल आँखें संक्रमण का संकेत देती हैं। साफ पानी से धोएं और दवा लगाएं।",
        "urgency": "routine",
        "follow_up_tests": ["Clinical eye examination", "Bacterial culture if severe"],
        "recommendations": ["Wash eyes 3-4 times daily", "Use antibiotic eye drops", "Avoid sharing towels"]
    },

    "hypercholesterolemia": {
        "disease_name": "High Cholesterol / Lipid Disorder",
        "visual_signs": {
            "corneal_arcus": 0.85,              # Gray/white ring at cornea edge
            "xanthelasma": 0.90,                # Yellow eyelid plaques
        },
        "clinical_rules": [
            {"condition": "corneal_arcus > 0.70 AND age < 40", "boost": 0.30},  # Early arcus = familial
            {"condition": "xanthelasma > 0.75 AND corneal_arcus > 0.70", "boost": 0.25},
        ],
        "explanation_en": "Eye deposits suggest high cholesterol. Get lipid panel test. Affects heart health.",
        "explanation_hi": "आँख में जमाव कोलेस्ट्रॉल का संकेत है। लिपिड पैनल परीक्षण करवाएं।",
        "urgency": "moderate",
        "follow_up_tests": ["Total Cholesterol", "LDL", "HDL", "Triglycerides", "Lipoprotein(a)"],
        "recommendations": ["Get lipid panel test", "Reduce saturated fat", "Increase exercise"]
    },

    "wilsons_disease": {
        "disease_name": "Wilson's Disease (Copper Overload)",
        "visual_signs": {
            "kayser_fleischer_ring": 0.95,     # Brown-green corneal ring (pathognomonic)
        },
        "clinical_rules": [
            {"condition": "kayser_fleischer_ring > 0.85", "boost": 0.50, "urgency_upgrade": "urgent"},
        ],
        "explanation_en": "URGENT: Distinctive eye sign suggests Wilson's disease (genetic disorder). See specialist immediately.",
        "explanation_hi": "तुरंत: आँख का विशिष्ट संकेत Wilson's disease दर्शाता है। विशेषज्ञ से तुरंत मिलें।",
        "urgency": "urgent",
        "follow_up_tests": ["Serum Copper", "Ceruloplasmin (Low)", "24-hour Urine Copper", "Slit-lamp exam"],
        "recommendations": ["URGENT: Hepatology specialist referral", "Copper chelation therapy needed"]
    },
}
class LLMEdgeCaseHandler:
    """
    Lightweight LLM handler for edge cases only.
    Uses Phi-3-mini or similar small model for unclear findings.
    """

    def __init__(self, use_llm=False, model_name="phi-3-mini"):
        self.use_llm = use_llm
        self.model_name = model_name
        self.llm_model = None

        if self.use_llm:
            try:
                # Try to load small model - adjust based on your setup
                from transformers import pipeline
                self.llm_model = pipeline(
                    "text-generation",
                    model=model_name,
                    device=0 if torch.cuda.is_available() else -1,
                    max_length=200,
                    top_p=0.95,
                )
                print(f"✓ LLM model {model_name} loaded for edge cases")
            except Exception as e:
                print(f"⚠ LLM not available: {e}. Using template-only mode.")
                self.use_llm = False

    def handle_edge_case(self, findings: str, age: int, language: str = "hi") -> str:
        """
        Use LLM ONLY for unclear/multiple disease cases.

        Args:
            findings: Description like "Patient shows pallor + yellow eyes + red palms"
            age: Patient age
            language: "en" or "hi"

        Returns:
            Simple 2-sentence clinical explanation
        """

        if not self.use_llm or not self.llm_model:
            return self._template_fallback(findings, language)

        if language == "hi":
            prompt = (
                f"रोगी की आयु {age} वर्ष है। "
                f"लक्षण: {findings}। "
                f"2 वाक्यों में सरल हिंदी में बताएं कि यह क्या रोग हो सकता है। "
                f"केवल चिकित्सा तथ्य दें।"
            )
        else:
            prompt = (
                f"Patient age {age}. "
                f"Findings: {findings}. "
                f"Explain in 2 sentences what this could indicate. Keep simple."
            )

        try:
            response = self.llm_model(prompt, max_new_tokens=100)
            explanation = response[0]["generated_text"].split(prompt)[-1].strip()
            return explanation[:200]  # Limit to 200 chars
        except Exception as e:
            print(f"LLM error: {e}. Falling back to templates.")
            return self._template_fallback(findings, language)

    def _template_fallback(self, findings: str, language: str) -> str:
        """Fallback when LLM unavailable"""
        if language == "hi":
            return "कई संकेत दिख रहे हैं। विशेषज्ञ से तुरंत जांच करवाएं।"
        else:
            return "Multiple findings detected. Please consult a specialist immediately."



class MedicalTemplatesEngine:
    """Analyzes CNN outputs and generates medical assessments."""

    def __init__(self, use_llm=False, model_name="phi-3-mini"):
        self.diseases = DISEASE_TEMPLATES
        self.llm_handler = LLMEdgeCaseHandler(use_llm=use_llm, model_name=llm_model_name)
        self.use_llm = use_llm
    def analyze(self, cnn_output, age=50, language="en"):
        """
        Main API for your app.
        
        Args:
            cnn_output: dict of {"sign_name": confidence_score, ...}
                       from your CNN model
            age: patient age
            language: "en" or "hi"
        
        Returns:
            list of predictions sorted by confidence
        """
        results = []
        
        for disease_id, template in self.diseases.items():
            score = self._score_disease(cnn_output, age, template)
            if score is not None:
                results.append({
                    "disease_id": disease_id,
                    "disease_name": template["disease_name"],
                    "confidence": score,
                    "urgency": self._get_urgency(cnn_output, age, template),
                    "explanation": template[f"explanation_{language}"],
                    "follow_up_tests": template["follow_up_tests"],
                    "recommendations": template["recommendations"],
                })
        
        # Sort by confidence
        return sorted(results, key=lambda x: x["confidence"], reverse=True)

    def _score_disease(self, cnn_output, age, template):
        """Score a disease based on visual signs."""
        base_score = 0.0
        sign_count = 0
        
        for sign, weight in template["visual_signs"].items():
            if sign in cnn_output:
                base_score += cnn_output[sign] * weight
                sign_count += 1
        
        if sign_count == 0:
            return None
        
        score = base_score / sign_count
        
        # Apply clinical rules
        for rule in template.get("clinical_rules", []):
            if self._check_rule(rule["condition"], cnn_output, age):
                score += rule.get("boost", 0)
        
        # Cap at 0.99, min 0.20
        score = max(0.0, min(score, 0.99))
        return round(score, 2) if score >= 0.20 else None

    def _get_urgency(self, cnn_output, age, template):
        """Determine urgency level based on rules."""
        urgency = template.get("urgency", "routine")
        
        for rule in template.get("clinical_rules", []):
            if self._check_rule(rule["condition"], cnn_output, age):
                if "urgency_upgrade" in rule:
                    urgency = rule["urgency_upgrade"]
        
        return urgency

    def _check_rule(self, condition, cnn_output, age):
        """Evaluate a condition string."""
        # Replace age placeholder
        condition = condition.replace("age", str(age))
        
        try:
            if " AND " in condition:
                return all(self._eval(part.strip(), cnn_output) for part in condition.split(" AND "))
            elif " OR " in condition:
                return any(self._eval(part.strip(), cnn_output) for part in condition.split(" OR "))
            return self._eval(condition, cnn_output)
        except:
            return False

    def _eval(self, cond, cnn_output):
        """Evaluate single comparison: 'sign > 0.75'"""
        for op in [">=", "<=", ">", "<"]:  # Check >= before >
            if op in cond:
                parts = cond.split(op)
                left = parts[0].strip()
                right = float(parts[1].strip())
                value = cnn_output.get(left, 0.0)
                
                if op == ">=" and value >= right:
                    return True
                elif op == "<=" and value <= right:
                    return True
                elif op == ">" and value > right:
                    return True
                elif op == "<" and value < right:
                    return True
        
        return False


# ============================================================================
# USAGE EXAMPLE
# ============================================================================

if __name__ == "__main__":
    engine = MedicalTemplatesEngine()
    
    # Simulate CNN output
    cnn_output = {
        "conjunctival_pallor": 0.78,
        "oral_mucosal_pallor": 0.72,
        "scleral_icterus": 0.12,
        "facial_asymmetry": 0.15,
    }
    
    results = engine.analyze(cnn_output, age=55, language="en")
    
    for i, result in enumerate(results[:3], 1):
        print(f"{i}. {result['disease_name']} ({result['urgency'].upper()})")
        print(f"   Confidence: {result['confidence']*100:.0f}%")
        print(f"   {result['explanation']}\n")