src/ai_processor.py

import os
import logging
import cv2
import numpy as np
from PIL import Image
import torch
import json
from datetime import datetime
import tensorflow as tf
from transformers import pipeline
from ultralytics import YOLO
from tensorflow.keras.models import load_model
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from huggingface_hub import HfApi, HfFolder
import spaces

from .config import Config

class AIProcessor:
    def __init__(self):
        self.models_cache = {}
        self.knowledge_base_cache = {}
        self.config = Config()
        self.px_per_cm = self.config.PIXELS_PER_CM
        self._initialize_models()

    def _initialize_models(self):
        """Initialize all AI models including real-time models"""
        try:
            # Set HuggingFace token
            if self.config.HF_TOKEN:
                HfFolder.save_token(self.config.HF_TOKEN)
                logging.info("HuggingFace token set successfully")

            # Initialize MedGemma pipeline for medical text generation
            try:
                self.models_cache["medgemma_pipe"] = pipeline(
                    "image-text-to-text",
                    model="google/medgemma-4b-it",
                    torch_dtype=torch.bfloat16,
                    offload_folder="offload",
                    device_map="cuda",
                    token=self.config.HF_TOKEN
                )
                logging.info("✅ MedGemma pipeline loaded successfully")
            except Exception as e:
                logging.warning(f"MedGemma pipeline not available: {e}")

            # Initialize YOLO model for wound detection
            try:
                self.models_cache["det"] = YOLO(self.config.YOLO_MODEL_PATH)
                logging.info("✅ YOLO detection model loaded successfully")
            except Exception as e:
                logging.warning(f"YOLO model not available: {e}")

            # Initialize segmentation model
            try:
                self.models_cache["seg"] = load_model(self.config.SEG_MODEL_PATH, compile=False)
                logging.info("✅ Segmentation model loaded successfully")
            except Exception as e:
                logging.warning(f"Segmentation model not available: {e}")

            # Initialize wound classification model
            try:
                self.models_cache["cls"] = pipeline(
                    "image-classification", 
                    model="Hemg/Wound-classification", 
                    token=self.config.HF_TOKEN, 
                    device="cpu"
                )
                logging.info("✅ Wound classification model loaded successfully")
            except Exception as e:
                logging.warning(f"Wound classification model not available: {e}")

            # Initialize embedding model for knowledge base
            try:
                self.models_cache["embedding_model"] = HuggingFaceEmbeddings(
                    model_name="sentence-transformers/all-MiniLM-L6-v2", 
                    model_kwargs={'device': 'cpu'}
                )
                logging.info("✅ Embedding model loaded successfully")
            except Exception as e:
                logging.warning(f"Embedding model not available: {e}")

            logging.info("✅ All models loaded.")
            self._load_knowledge_base()

        except Exception as e:
            logging.error(f"Error initializing AI models: {e}")

    def _load_knowledge_base(self):
        """Load knowledge base from PDF guidelines"""
        try:
            documents = []
            for pdf_path in self.config.GUIDELINE_PDFS:
                if os.path.exists(pdf_path):
                    loader = PyPDFLoader(pdf_path)
                    docs = loader.load()
                    documents.extend(docs)
                    logging.info(f"Loaded PDF: {pdf_path}")

            if documents and 'embedding_model' in self.models_cache:
                # Split documents into chunks
                text_splitter = RecursiveCharacterTextSplitter(
                    chunk_size=1000,
                    chunk_overlap=100
                )
                chunks = text_splitter.split_documents(documents)
                
                # Create vector store
                vectorstore = FAISS.from_documents(chunks, self.models_cache['embedding_model'])
                self.knowledge_base_cache['vectorstore'] = vectorstore
                logging.info(f"✅ Knowledge base loaded with {len(chunks)} chunks")
            else:
                self.knowledge_base_cache['vectorstore'] = None
                logging.warning("Knowledge base not available - no PDFs found or embedding model unavailable")

        except Exception as e:
            logging.warning(f"Knowledge base loading error: {e}")
            self.knowledge_base_cache['vectorstore'] = None

    @spaces.GPU(enable_queue=True, duration=120) 
    def perform_visual_analysis(self, image_pil):
        """Perform comprehensive visual analysis of wound image."""
        try:
            # Convert PIL to OpenCV format
            image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    
            # YOLO detection
            if 'det' not in self.models_cache:
                raise ValueError("YOLO detection model not available.")
    
            results = self.models_cache['det'].predict(image_cv, verbose=False, device="cpu")
    
            if not results or not results[0].boxes:
                raise ValueError("No wound detected in the image.")
    
            # Extract bounding box
            box = results[0].boxes[0].xyxy[0].cpu().numpy().astype(int)
            x1, y1, x2, y2 = box
            region_cv = image_cv[y1:y2, x1:x2]
    
            # Save detection image
            detection_image_cv = image_cv.copy()
            cv2.rectangle(detection_image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
            os.makedirs(os.path.join(self.config.UPLOADS_DIR, "analysis"), exist_ok=True)
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            detection_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"detection_{timestamp}.png")
            cv2.imwrite(detection_image_path, detection_image_cv)
            detection_image_pil = Image.fromarray(cv2.cvtColor(detection_image_cv, cv2.COLOR_BGR2RGB))
    
            # Initialize outputs
            length = breadth = area = 0
            segmentation_image_pil = None
            segmentation_image_path = None
    
            # Segmentation (optional)
            if 'seg' in self.models_cache:
                input_size = self.models_cache['seg'].input_shape[1:3]  # (height, width)
                resized_region = cv2.resize(region_cv, (input_size[1], input_size[0]))
    
                seg_input = np.expand_dims(resized_region / 255.0, 0)
                mask_pred = self.models_cache['seg'].predict(seg_input, verbose=0)[0]
                mask_np = (mask_pred[:, :, 0] > 0.5).astype(np.uint8)
    
                # Resize mask back to original region size
                mask_resized = cv2.resize(mask_np, (region_cv.shape[1], region_cv.shape[0]), interpolation=cv2.INTER_NEAREST)
    
                # Overlay mask on region for visualization
                overlay = region_cv.copy()
                overlay[mask_resized == 1] = [0, 0, 255]  # Red overlay
    
                # Blend overlay for final output
                segmented_visual = cv2.addWeighted(region_cv, 0.7, overlay, 0.3, 0)
    
                # Save segmentation image
                segmentation_image_path = os.path.join(self.config.UPLOADS_DIR, "analysis", f"segmentation_{timestamp}.png")
                cv2.imwrite(segmentation_image_path, segmented_visual)
                segmentation_image_pil = Image.fromarray(cv2.cvtColor(segmented_visual, cv2.COLOR_BGR2RGB))
    
                # Wound measurements from resized mask
                contours, _ = cv2.findContours(mask_resized, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                if contours:
                    cnt = max(contours, key=cv2.contourArea)
                    x, y, w, h = cv2.boundingRect(cnt)
                    length = round(h / self.px_per_cm, 2)
                    breadth = round(w / self.px_per_cm, 2)
                    area = round(cv2.contourArea(cnt) / (self.px_per_cm ** 2), 2)
    
            # Classification (optional)
            wound_type = "Unknown"
            if 'cls' in self.models_cache:
                try:
                    region_pil = Image.fromarray(cv2.cvtColor(region_cv, cv2.COLOR_BGR2RGB))
                    cls_result = self.models_cache['cls'](region_pil)
                    wound_type = max(cls_result, key=lambda x: x['score'])['label']
                except Exception as e:
                    logging.warning(f"Wound classification error: {e}")
    
            return {
                'wound_type': wound_type,
                'length_cm': length,
                'breadth_cm': breadth,
                'surface_area_cm2': area,
                'detection_confidence': float(results[0].boxes[0].conf.cpu().item()),
                'bounding_box': box.tolist(),
                'detection_image_path': detection_image_path,
                'detection_image_pil': detection_image_pil,
                'segmentation_image_path': segmentation_image_path,
                'segmentation_image_pil': segmentation_image_pil
            }
    
        except Exception as e:
            logging.error(f"Visual analysis error: {e}")
            raise ValueError(f"Visual analysis failed: {str(e)}")


    def query_guidelines(self, query: str):
        """Query the knowledge base for relevant guidelines"""
        try:
            vector_store = self.knowledge_base_cache.get("vectorstore")
            if not vector_store:
                return "Knowledge base unavailable - clinical guidelines not loaded"

            # Retrieve relevant documents
            retriever = vector_store.as_retriever(search_kwargs={"k": 10})
            docs = retriever.invoke(query)
            
            if not docs:
                return "No relevant guidelines found for the query"

            # Format the results
            formatted_results = []
            for doc in docs:
                source = doc.metadata.get('source', 'Unknown')
                page = doc.metadata.get('page', 'N/A')
                content = doc.page_content.strip()
                formatted_results.append(f"Source: {source}, Page: {page}\nContent: {content}")

            return "\n\n".join(formatted_results)

        except Exception as e:
            logging.error(f"Guidelines query error: {e}")
            return f"Error querying guidelines: {str(e)}"

    @spaces.GPU(enable_queue=True, duration=120) 
    def generate_final_report(self, patient_info, visual_results, guideline_context, image_pil, max_new_tokens=None):
        """Generate comprehensive medical report using MedGemma"""
        try:
            if 'medgemma_pipe' not in self.models_cache:
                return self._generate_fallback_report(patient_info, visual_results, guideline_context)

            max_tokens = max_new_tokens or self.config.MAX_NEW_TOKENS
            
            # Get detection and segmentation images if available
            detection_image = visual_results.get('detection_image_pil', None)
            segmentation_image = visual_results.get('segmentation_image_pil', None)
            
            # Create image paths for report
            detection_path = visual_results.get('detection_image_path', '')
            segmentation_path = visual_results.get('segmentation_image_path', '')
            
            # Create detailed prompt for medical analysis with image paths
            prompt = f"""
# Wound Care Report

## Patient Information
{patient_info}

## Visual Analysis Summary
- Wound Type: {visual_results.get('wound_type', 'Unknown')}
- Length: {visual_results.get('length_cm', 0)} cm
- Breadth: {visual_results.get('breadth_cm', 0)} cm
- Surface Area: {visual_results.get('surface_area_cm2', 0)} cm²
- Detection Confidence: {visual_results.get('detection_confidence', 0):.2f}

## Clinical Reference
{guideline_context}

You are SmartHeal-AI Agent, a world-class wound care AI specialist trained in clinical wound assessment and guideline-based treatment planning.
Your task is to process the following structured inputs (patient data, wound measurements, clinical guidelines, and image) and perform **clinical reasoning and decision-making** to generate a complete wound care report.
---
🔍 **YOUR PROCESS — FOLLOW STRICTLY:**
### Step 1: Clinical Reasoning (Chain-of-Thought)
Use the provided information to think step-by-step about:
- Patient’s risk factors (e.g. diabetes, age, healing limitations)
- Wound characteristics (size, tissue appearance, moisture, infection signs)
- Visual clues from the image (location, granulation, maceration, inflammation, surrounding skin)

---
-Step 2: Structured Clinical Report
Generate the following report sections using markdown and medical terminology:
**1. Clinical Summary**
- Describe wound appearance and tissue types (e.g., slough, necrotic, granulating, epithelializing)
- Include size, wound bed condition, peri-wound skin, and signs of infection or biofilm
- Mention inferred location (e.g., heel, forefoot) if image allows
- Summarize patient's systemic risk profile
**2. Medicinal & Dressing Recommendations**
Based on your analysis:
- Recommend specific **wound care dressings** (e.g., hydrocolloid, alginate, foam, antimicrobial silver, etc.) suitable to wound moisture level and infection risk
- Propose **topical or systemic agents** ONLY if relevant — include name classes (e.g., antiseptic: povidone iodine, antibiotic ointments, enzymatic debriders)
- Mention **techniques** (e.g., sharp debridement, NPWT, moisture balance, pressure offloading, dressing frequency)
- Avoid repeating guidelines — **apply them**
**3. Key Risk Factors**
Explain how the patient’s condition (e.g., diabetic, poor circulation, advanced age, poor hygiene) may affect wound healing
**4. Prognosis & Monitoring Advice**
- Mention how often wound should be reassessed
- Indicate signs to monitor for deterioration or improvement
- Include when escalation to specialist is necessary

**Note:** Every dressing change is a chance for wound reassessment. Always perform a thorough wound evaluation at each dressing change.
"""

            # Prepare messages for MedGemma with all available images
            content_list = [{"type": "text", "text": prompt}]
            
            # Add original image
            if image_pil:
                content_list.insert(0, {"type": "image", "image": image_pil})
                
            # Add detection image if available
            if detection_image:
                content_list.insert(1, {"type": "image", "image": detection_image})
                
            # Add segmentation image if available
            if segmentation_image:
                content_list.insert(2, {"type": "image", "image": segmentation_image})
            
            messages = [
                {
                    "role": "system",
                    "content": [{"type": "text", "text": "You are a world-class medical AI assistant specializing in wound care with expertise in wound assessment and treatment. Provide concise, evidence-based medical assessments focusing on: (1) Precise wound classification based on tissue type and appearance, (2) Specific treatment recommendations with exact product names or interventions when appropriate, (3) Objective evaluation of healing progression or deterioration indicators, and (4) Clear follow-up timelines. Avoid general statements and prioritize actionable insights based on the visual analysis measurements and patient context."}],
                },
                {
                    "role": "user",
                    "content": content_list
                }
            ]

            # Generate report using MedGemma
            output = self.models_cache['medgemma_pipe'](
                text=messages,
                max_new_tokens=1024,
                do_sample=False,
            )
            
            generated_content = output[0]['generated_text'][-1].get('content', '').strip()
            
            # Include image paths in the final report for display in UI
            if generated_content:
                # Add image paths to the report for frontend display
                image_paths_section = f"""
## Analysis Images
- Original Image: {image_pil}
- Detection Image: {detection_path}
- Segmentation Image: {segmentation_path}
"""
                generated_content = image_paths_section + generated_content
                
            return generated_content if generated_content else self._generate_fallback_report(patient_info, visual_results, guideline_context)

        except Exception as e:
            logging.error(f"MedGemma report generation error: {e}")
            return self._generate_fallback_report(patient_info, visual_results, guideline_context)

    def _generate_fallback_report(self, patient_info, visual_results, guideline_context):
        """Generate a fallback report when MedGemma is not available"""
        # Get image paths for report
        detection_path = visual_results.get('detection_image_path', 'Not available')
        segmentation_path = visual_results.get('segmentation_image_path', 'Not available')
        
        report = f"""
# Wound Analysis Report
## Patient Information
{patient_info}

## Visual Analysis Results
- **Wound Type**: {visual_results.get('wound_type', 'Unknown')}
- **Dimensions**: {visual_results.get('length_cm', 0)} cm × {visual_results.get('breadth_cm', 0)} cm
- **Surface Area**: {visual_results.get('surface_area_cm2', 0)} cm²
- **Detection Confidence**: {visual_results.get('detection_confidence', 0):.2f}

## Analysis Images
- **Detection Image**: {detection_path}
- **Segmentation Image**: {segmentation_path}

## Assessment
Based on the visual analysis, this appears to be a {visual_results.get('wound_type', 'wound')} with measurable dimensions. 

## Recommendations
- Continue monitoring wound healing progress
- Maintain proper wound hygiene
- Follow appropriate dressing protocols
- Seek medical attention if signs of infection develop

## Clinical Guidelines
{guideline_context[:500]}...

*Note: This is an automated analysis. Please consult with a healthcare professional for definitive diagnosis and treatment.*
"""
        return report

    def save_and_commit_image(self, image_pil):
        """Save image locally and optionally upload to HuggingFace dataset"""
        try:
            # Ensure uploads directory exists
            os.makedirs(self.config.UPLOADS_DIR, exist_ok=True)
            
            # Generate filename with timestamp
            filename = f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.png"
            local_path = os.path.join(self.config.UPLOADS_DIR, filename)
            
            # Save image locally
            image_pil.convert("RGB").save(local_path)
            logging.info(f"Image saved locally: {local_path}")

            # Upload to HuggingFace dataset if configured
            if self.config.HF_TOKEN and self.config.DATASET_ID:
                try:
                    api = HfApi()
                    api.upload_file(
                        path_or_fileobj=local_path,
                        path_in_repo=f"images/{filename}",
                        repo_id=self.config.DATASET_ID,
                        repo_type="dataset",
                        commit_message=f"Upload wound image: {filename}"
                    )
                    logging.info("✅ Image uploaded to HuggingFace dataset")
                except Exception as e:
                    logging.warning(f"HuggingFace upload failed: {e}")

            return local_path

        except Exception as e:
            logging.error(f"Image saving error: {e}")
            return None

    
    def full_analysis_pipeline(self, image, questionnaire_data):
        """Complete analysis pipeline with real-time models"""
        try:
            # Save the image
            saved_path = self.save_and_commit_image(image)
            
            # Perform visual analysis
            visual_results = self.perform_visual_analysis(image)
            
            # Format patient information
            patient_info = ", ".join([f"{k}: {v}" for k, v in questionnaire_data.items() if v])
            
            # Create query for guidelines
            wound_type = visual_results.get('wound_type', 'wound')
            moisture = questionnaire_data.get('moisture', 'unknown')
            infection = questionnaire_data.get('infection', 'unknown')
            diabetic = questionnaire_data.get('diabetic', 'unknown')
            
            query = f"best practices for managing a {wound_type} with moisture level '{moisture}' and signs of infection '{infection}' in a patient who is diabetic '{diabetic}'"
            
            # Query guidelines
            guideline_context = self.query_guidelines(query)
            
            # Generate final report
            final_report = self.generate_final_report(patient_info, visual_results, guideline_context, image)
            
            return {
                'success': True,
                'visual_analysis': visual_results,
                'report': final_report,
                'saved_image_path': saved_path,
                'timestamp': datetime.now().isoformat()
            }

        except Exception as e:
            logging.error(f"Full analysis pipeline error: {e}")
            return {
                'success': False,
                'error': str(e),
                'timestamp': datetime.now().isoformat()
            }

    # Legacy methods for backward compatibility
    def analyze_wound(self, image, questionnaire_data):
        """Legacy method for backward compatibility"""
        try:
            # Convert string path to PIL Image if needed
            if isinstance(image, str):
                try:
                    from PIL import Image
                    image = Image.open(image)
                    logging.info(f"Converted string path to PIL Image: {image}")
                except Exception as e:
                    logging.error(f"Error converting string path to image: {e}")
            
            # Ensure we have a PIL Image object
            if not isinstance(image, Image.Image):
                try:
                    from PIL import Image
                    import io
                    
                    # If it's a file-like object
                    if hasattr(image, 'read'):
                        # Reset file pointer if possible
                        if hasattr(image, 'seek'):
                            image.seek(0)
                        image = Image.open(image)
                        logging.info("Converted file-like object to PIL Image")
                except Exception as e:
                    logging.error(f"Error ensuring image is PIL Image: {e}")
                    raise ValueError(f"Invalid image format: {type(image)}")
            
            result = self.full_analysis_pipeline(image, questionnaire_data)
            
            if result['success']:
                return {
                    'timestamp': result['timestamp'],
                    'summary': f"Analysis completed for {questionnaire_data.get('patient_name', 'patient')}",
                    'recommendations': result['report'],
                    'wound_detection': {
                        'status': 'success',
                        'detections': [result['visual_analysis']],
                        'total_wounds': 1
                    },
                    'segmentation_result': {
                        'status': 'success',
                        'wound_area_percentage': result['visual_analysis'].get('surface_area_cm2', 0)
                    },
                    'risk_assessment': self._assess_risk_legacy(questionnaire_data),
                    'guideline_recommendations': [result['report'][:200] + "..."]
                }
            else:
                return {
                    'timestamp': result['timestamp'],
                    'summary': f"Analysis failed: {result['error']}",
                    'recommendations': "Please consult with a healthcare professional.",
                    'wound_detection': {'status': 'error', 'message': result['error']},
                    'segmentation_result': {'status': 'error', 'message': result['error']},
                    'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
                    'guideline_recommendations': ["Analysis unavailable due to error"]
                }
                
        except Exception as e:
            logging.error(f"Legacy analyze_wound error: {e}")
            return {
                'timestamp': datetime.now().isoformat(),
                'summary': f"Analysis error: {str(e)}",
                'recommendations': "Please consult with a healthcare professional.",
                'wound_detection': {'status': 'error', 'message': str(e)},
                'segmentation_result': {'status': 'error', 'message': str(e)},
                'risk_assessment': {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []},
                'guideline_recommendations': ["Analysis unavailable due to error"]
            }

    def _assess_risk_legacy(self, questionnaire_data):
        """Legacy risk assessment for backward compatibility"""
        risk_factors = []
        risk_score = 0
        
        try:
            # Age factor
            age = questionnaire_data.get('patient_age', 0)
            if age > 65:
                risk_factors.append("Advanced age (>65)")
                risk_score += 2
            elif age > 50:
                risk_factors.append("Older adult (50-65)")
                risk_score += 1
            
            # Duration factor
            duration = questionnaire_data.get('wound_duration', '').lower()
            if any(term in duration for term in ['month', 'months', 'year']):
                risk_factors.append("Chronic wound (>4 weeks)")
                risk_score += 3
            
            # Pain level
            pain_level = questionnaire_data.get('pain_level', 0)
            if pain_level >= 7:
                risk_factors.append("High pain level")
                risk_score += 2
            
            # Medical history risk factors
            medical_history = questionnaire_data.get('medical_history', '').lower()
            if 'diabetes' in medical_history:
                risk_factors.append("Diabetes mellitus")
                risk_score += 3
            if 'circulation' in medical_history or 'vascular' in medical_history:
                risk_factors.append("Vascular/circulation issues")
                risk_score += 2
            if 'immune' in medical_history:
                risk_factors.append("Immune system compromise")
                risk_score += 2
            
            # Determine risk level
            if risk_score >= 7:
                risk_level = "High"
            elif risk_score >= 4:
                risk_level = "Moderate"
            else:
                risk_level = "Low"
            
            return {
                'risk_score': risk_score,
                'risk_level': risk_level,
                'risk_factors': risk_factors
            }
            
        except Exception as e:
            logging.error(f"Risk assessment error: {e}")
            return {'risk_score': 0, 'risk_level': 'Unknown', 'risk_factors': []}