"""
battlefield_triage_ai.py
MARCH/PAWS-compliant AI triage with explainable decisions
Built on TCCC guidelines and real-world SF medical experience
Date: 2025-05-10
Author: Jeremy Cleland
"""

import pandas as pd
import numpy as np
from datetime import datetime
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import shap
import logging

# Set up battlefield-appropriate logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

class BattlefieldTriageAI:
    """
    AI-enhanced MARCH/PAWS triage following SF medical protocols

    This system augments—never replaces—medic judgment. Final
    decisions always rest with the human on scene.
    """

    def __init__(self):
        self.model = RandomForestClassifier(
            n_estimators=100,
            class_weight='balanced',
            random_state=42,
            n_jobs=-1  # Use all available cores
        )
        self.scaler = StandardScaler()
        self.internal_hemorrhage_model = RandomForestClassifier(
            n_estimators=50,
            random_state=42
        )
        logger.info("Battlefield Triage AI initialized")

    def calculate_shock_index(self, heart_rate: float, systolic_bp: float) -> float:
        """
        Shock Index = HR/SBP
        > 0.9: Increased mortality risk
        > 1.0: Significant hemorrhage likely

        Reference: Ranger Medic Handbook (2022)
        """
        if systolic_bp <= 0:
            logger.warning("Invalid SBP detected, returning max shock index")
            return 2.0  # Max value indicating critical state
        return heart_rate / systolic_bp

    def assess_hemorrhage_risk(self, vitals: dict) -> np.ndarray:
        """
        Detect subtle signs of internal hemorrhage:
        - Narrowing pulse pressure (<30 mmHg)
        - Elevated shock index (>1.0)
        - Mechanism suggesting high-energy trauma

        Based on TCCC internal hemorrhage protocols
        """
        features = []

        # Pulse pressure analysis
        pulse_pressure = vitals['systolic_bp'] - vitals['diastolic_bp']
        features.append(1 if pulse_pressure < 30 else 0)

        # Shock index calculation
        shock_index = self.calculate_shock_index(
            vitals['heart_rate'],
            vitals['systolic_bp']
        )
        features.append(1 if shock_index > 1.0 else 0)

        # Mechanism of injury risk
        moi_risk = vitals.get('blast_injury', 0)
        features.append(moi_risk)

        return np.array(features)

    def extract_march_features(self, patient_data: dict) -> dict:
        """
        Convert battlefield observations to AI features
        Aligned with MARCH/PAWS assessment sequence
        """
        # Core calculations
        shock_index = self.calculate_shock_index(
            patient_data['heart_rate'],
            patient_data['systolic_bp']
        )
        pulse_pressure = patient_data['systolic_bp'] - patient_data['diastolic_bp']
        map_value = (patient_data['systolic_bp'] + 2 * patient_data['diastolic_bp']) / 3

        features = {
            # M - Massive Hemorrhage
            'shock_index': shock_index,
            'pulse_pressure': pulse_pressure,
            'tourniquet_applied': patient_data.get('tourniquet_applied', 0),
            'external_bleeding_controlled': patient_data.get('bleeding_controlled', 0),

            # A - Airway
            'spo2': patient_data.get('spo2', 98),
            'airway_patent': patient_data.get('airway_patent', 1),
            'stridor_present': patient_data.get('stridor', 0),

            # R - Respiration
            'respiratory_rate': patient_data['respiratory_rate'],
            'breath_sounds_equal': patient_data.get('equal_breath_sounds', 1),
            'chest_decompressed': patient_data.get('ncd_performed', 0),

            # C - Circulation
            'mean_arterial_pressure': map_value,
            'capillary_refill': patient_data.get('cap_refill', 2),
            'txa_administered': patient_data.get('txa_given', 0),

            # H - Head injury/Hypothermia
            'gcs_score': patient_data.get('gcs', 15),
            'pupil_response': patient_data.get('pupils_equal', 1),
            'core_temperature': patient_data.get('temperature', 37.0),

            # Additional risk factors
            'mechanism_blast': patient_data.get('blast_injury', 0),
            'time_since_injury': patient_data.get('injury_time_minutes', 0)
        }

        # Calculate composite risk scores
        features['hemorrhage_risk_score'] = self.assess_hemorrhage_risk(patient_data).sum()

        return features

    def predict_priority(self, patient_vitals: dict) -> dict:
        """
        Generate triage priority with confidence metrics
        Returns structured assessment following MARCH sequence
        """
        features = self.extract_march_features(patient_vitals)
        features_df = pd.DataFrame([features])

        # Standardize features
        features_scaled = self.scaler.transform(features_df)

        # Get model predictions
        priority = self.model.predict(features_scaled)[0]
        priority_proba = self.model.predict_proba(features_scaled)[0]

        # Assess internal hemorrhage risk separately
        hemorrhage_features = self.assess_hemorrhage_risk(patient_vitals)
        internal_risk = self.internal_hemorrhage_model.predict_proba(
            hemorrhage_features.reshape(1, -1)
        )[0][1]

        # Generate comprehensive assessment
        assessment = self._generate_march_report(features, priority, patient_vitals)

        return {
            'triage_category': self._map_priority_to_category(priority),
            'priority_score': priority,
            'confidence': float(max(priority_proba)),
            'internal_hemorrhage_risk': float(internal_risk),
            'march_assessment': assessment,
            'recommended_actions': self._get_immediate_actions(features, priority),
            'evacuation_priority': self._determine_evac_priority(features, priority)
        }

    def _generate_march_report(self, features: dict, priority: int,
                               raw_vitals: dict) -> list:
        """
        Generate structured MARCH assessment
        Mimics how an SF medic would report up the chain
        """
        report = []

        # M - Massive Hemorrhage
        if features['shock_index'] > 1.0:
            report.append("M: CRITICAL - Shock index >1.0, massive hemorrhage likely")
        elif features['pulse_pressure'] < 30:
            report.append("M: WARNING - Narrow pulse pressure, occult bleeding possible")
        elif not features['external_bleeding_controlled']:
            report.append("M: URGENT - External hemorrhage not yet controlled")
        else:
            report.append("M: STABLE - No active hemorrhage detected")

        # A - Airway
        if features['spo2'] < 90:
            report.append("A: CRITICAL - Severe hypoxia, immediate airway intervention")
        elif features['stridor_present']:
            report.append("A: URGENT - Stridor present, airway compromise developing")
        elif not features['airway_patent']:
            report.append("A: WARNING - Airway requires continuous monitoring")
        else:
            report.append("A: CLEAR - Airway patent and protected")

        # R - Respiration
        rr = features['respiratory_rate']
        if rr > 29 or rr < 10:
            report.append(f"R: CRITICAL - RR {rr}, consider immediate NCD if indicated")
        elif not features['breath_sounds_equal']:
            report.append("R: WARNING - Unequal breath sounds, monitor for pneumothorax")
        else:
            report.append("R: ADEQUATE - Normal respiratory pattern")

        # C - Circulation
        map_val = features['mean_arterial_pressure']
        if map_val < 65:
            report.append(f"C: CRITICAL - MAP {map_val:.0f}, initiate damage control resuscitation")
        elif features['capillary_refill'] > 3:
            report.append("C: WARNING - Delayed cap refill, early shock state")
        else:
            report.append("C: STABLE - Hemodynamically stable")

        # H - Head injury/Hypothermia
        gcs = features['gcs_score']
        temp = features['core_temperature']
        if gcs < 15:
            report.append(f"H: WARNING - GCS {gcs}, monitor for deterioration")
        if temp < 36:
            report.append(f"H: HYPOTHERMIC - Core temp {temp}°C, initiate warming")

        # Time-critical factors
        injury_time = features.get('time_since_injury', 0)
        if injury_time > 60:
            report.append(f"TIME: {injury_time} minutes since injury - Golden Hour expiring")

        return report

    def _map_priority_to_category(self, priority: int) -> str:
        """Map numeric priority to NATO triage categories"""
        categories = {
            1: "IMMEDIATE (T1)",
            2: "DELAYED (T2)",
            3: "MINIMAL (T3)",
            4: "EXPECTANT (T4)"
        }
        return categories.get(priority, "IMMEDIATE (T1)")

    def _get_immediate_actions(self, features: dict, priority: int) -> list:
        """
        Generate immediate action items based on MARCH assessment
        These augment, not replace, medic judgment
        """
        actions = []

        # Hemorrhage control
        if features['shock_index'] > 1.0 and not features['tourniquet_applied']:
            actions.append("Apply tourniquet to suspected hemorrhage source")
        if features['hemorrhage_risk_score'] > 1 and not features['txa_administered']:
            actions.append("Administer TXA 2g IV/IO slow push")

        # Airway management
        if features['spo2'] < 90:
            actions.append("Secure airway - consider surgical cricothyroidotomy")
        elif features['stridor_present']:
            actions.append("Prepare for emergency airway - position for rapid intervention")

        # Respiratory support
        if features['respiratory_rate'] > 29 and not features['chest_decompressed']:
            actions.append("Assess for tension pneumothorax - NCD if indicated")

        # Circulatory support
        if features['mean_arterial_pressure'] < 65:
            actions.append("Initiate whole blood transfusion protocol")
            actions.append("Establish second large-bore IV/IO access")

        # Hypothermia prevention
        if features['core_temperature'] < 36:
            actions.append("Deploy HPMK - aggressive warming measures")

        return actions

    def _determine_evac_priority(self, features: dict, priority: int) -> str:
        """
        Determine MEDEVAC precedence based on injury pattern
        Reference: ATP 4-02.2 Medical Evacuation
        """
        if priority == 1 or features['shock_index'] > 1.5:
            return "URGENT - Evacuate within 1 hour"
        elif priority == 2 or features['hemorrhage_risk_score'] > 1:
            return "PRIORITY - Evacuate within 4 hours"
        elif priority == 3:
            return "ROUTINE - Evacuate within 24 hours"
        else:
            return "CONVENIENCE - Evacuate when possible"

    def explain_decision(self, patient_vitals: dict) -> dict:
        """
        Provide transparent explanation of AI reasoning
        Critical for building trust with medics in field
        """
        features = self.extract_march_features(patient_vitals)
        features_df = pd.DataFrame([features])

        # Create SHAP explainer
        explainer = shap.TreeExplainer(self.model)
        shap_values = explainer.shap_values(self.scaler.transform(features_df))

        # Handle multi-class output
        if isinstance(shap_values, list):
            shap_values = shap_values[1]  # Focus on high-priority class

        # Generate feature importance
        feature_importance = dict(zip(features_df.columns, shap_values[0]))

        # Create human-readable explanation
        top_factors = sorted(
            feature_importance.items(),
            key=lambda x: abs(x[1]),
            reverse=True
        )[:5]

        explanation = {
            'primary_drivers': [],
            'confidence_factors': [],
            'risk_indicators': []
        }

        for feature, importance in top_factors:
            if importance > 0:
                explanation['risk_indicators'].append(
                    f"{feature.replace('_', ' ').title()}: Increases priority (impact: {importance:.3f})"
                )
            else:
                explanation['confidence_factors'].append(
                    f"{feature.replace('_', ' ').title()}: Decreases priority (impact: {importance:.3f})"
                )

        # Add clinical context
        if features['shock_index'] > 1.0:
            explanation['primary_drivers'].append(
                "Shock index >1.0 strongly indicates ongoing hemorrhage"
            )
        if features['spo2'] < 90:
            explanation['primary_drivers'].append(
                "Critical hypoxia requires immediate intervention"
            )

        return explanation

# Demonstration with realistic combat casualty
if __name__ == "__main__":
    logger.info("Initializing Battlefield Triage AI system")
    ai_assistant = BattlefieldTriageAI()

    # Simulate IED blast casualty
    blast_casualty = {
        'heart_rate': 124,
        'systolic_bp': 86,
        'diastolic_bp': 52,
        'respiratory_rate': 28,
        'spo2': 89,
        'temperature': 35.5,
        'gcs': 13,
        'blast_injury': 1,
        'tourniquet_applied': 1,
        'txa_given': 0,
        'equal_breath_sounds': 0,
        'injury_time_minutes': 15
    }

    # Get AI assessment
    logger.info("Processing casualty assessment")
    result = ai_assistant.predict_priority(blast_casualty)

    print("\n" + "="*50)
    print("MARCH/PAWS AI TRIAGE ASSESSMENT")
    print("="*50)
    print(f"Time: {datetime.now().strftime('%H:%M:%S')}")
    print(f"Category: {result['triage_category']}")
    print(f"Confidence: {result['confidence']:.1%}")
    print(f"Internal Hemorrhage Risk: {result['internal_hemorrhage_risk']:.1%}")
    print(f"Evacuation: {result['evacuation_priority']}")

    print("\nMARCH Assessment:")
    for item in result['march_assessment']:
        print(f"  • {item}")

    print("\nImmediate Actions:")
    for action in result['recommended_actions']:
        print(f"  → {action}")

    # Get explanation
    explanation = ai_assistant.explain_decision(blast_casualty)
    print("\nAI Decision Factors:")
    for driver in explanation['primary_drivers']:
        print(f"  ⚠ {driver}")