""" Data Handler Module Handles all database operations and data persistence """ from datetime import datetime import pandas as pd import logging from database import get_db, test_db_connection from typing import Dict, List, Optional, Any # Configure logging logger = logging.getLogger(__name__) class DataHandler: """Handles all database operations for the virus prediction app""" def __init__(self): self.db = None self._initialize_db() def _initialize_db(self): """Initialize database connection""" try: self.db = get_db() if self.db is not None: # Create indexes for better performance self._create_indexes() logger.info("DataHandler initialized successfully") else: logger.warning("Failed to initialize database connection") except Exception as e: logger.error(f"Error initializing DataHandler: {e}") def _get_next_patient_id(self) -> str: """Generate auto-incrementing patient ID (P001, P002, etc.)""" try: if self.db is None: return "P001" # Get the counter collection for patient IDs counters = self.db['counters'] # Find and increment the patient counter result = counters.find_one_and_update( {'_id': 'patient_id'}, {'$inc': {'sequence_value': 1}}, upsert=True, return_document=True ) # Format as P001, P002, etc. sequence_num = result.get('sequence_value', 1) return f"P{sequence_num:03d}" except Exception as e: logger.error(f"Error generating patient ID: {e}") # Fallback to timestamp-based ID import time return f"P{int(time.time())}" def _create_indexes(self): """Create database indexes for better performance""" try: if self.db is None: return # Create indexes on frequently queried fields collections = { 'predictions': [ ('timestamp', -1), ('patient_id', 1), ('predicted_virus', 1), ('validation.validated', 1), ('validation.actual_virus_name', 1) ], 'patients': [ ('patient_id', 1), ('created_at', -1) ], 'usage_stats': [ ('date', -1), ('prediction_count', 1) ] } for collection_name, indexes in collections.items(): collection = self.db[collection_name] for index_fields in indexes: try: collection.create_index([index_fields]) except Exception as e: logger.warning(f"Index creation warning for {collection_name}: {e}") except Exception as e: logger.error(f"Error creating indexes: {e}") def save_prediction(self, patient_data: Dict, prediction_result: Dict, models_info: Dict = None, state_name: str = None, district_name: str = None) -> Optional[str]: """ Save prediction result to single collection with human-readable values Args: patient_data: Patient information and symptoms (encoded values) prediction_result: Model prediction results models_info: Model version and metadata state_name: Human-readable state name district_name: Human-readable district name Returns: Document ID if successful, None otherwise """ try: if self.db is None: logger.error("Database not initialized") return None # Use single collection for all data collection = self.db['virus_predictions'] # Generate unique patient ID patient_id = self._get_next_patient_id() # Transform patient data to human-readable format readable_patient_info = { 'patient_id': patient_id, # 'patient_name': patient_data.get('patient_name', ''), 'age': patient_data.get('age'), 'sex': 'Male' if patient_data.get('SEX') == 1 else 'Female', 'patient_type': 'Inpatient' if patient_data.get('PATIENTTYPE') == 1 else 'Outpatient', 'duration_of_illness_days': patient_data.get('durationofillness'), 'state_name': state_name or 'Unknown', 'district_name': district_name or 'Unknown', 'syndrome_name': patient_data.get('syndrome_name', ''), 'syndrome_specification': patient_data.get('other_syndrome_specification', ''), 'month_name': self._get_month_name(patient_data.get('month', 1)), 'year': patient_data.get('year') } # Transform symptoms to human-readable format (flat structure for CSV) symptoms_readable = self._transform_symptoms_to_readable(patient_data) # Transform prediction results to human-readable prediction_readable = { 'predicted_virus_name': prediction_result.get('predicted_virus'), 'prediction_confidence_percent': prediction_result.get('confidence'), 'top_1_virus': prediction_result.get('top_5_predictions', [{}])[0].get('virus', ''), 'top_1_confidence': prediction_result.get('top_5_predictions', [{}])[0].get('confidence', 0), 'top_2_virus': prediction_result.get('top_5_predictions', [{}])[1].get('virus', '') if len(prediction_result.get('top_5_predictions', [])) > 1 else '', 'top_2_confidence': prediction_result.get('top_5_predictions', [{}])[1].get('confidence', 0) if len(prediction_result.get('top_5_predictions', [])) > 1 else 0, 'top_3_virus': prediction_result.get('top_5_predictions', [{}])[2].get('virus', '') if len(prediction_result.get('top_5_predictions', [])) > 2 else '', 'top_3_confidence': prediction_result.get('top_5_predictions', [{}])[2].get('confidence', 0) if len(prediction_result.get('top_5_predictions', [])) > 2 else 0, 'top_4_virus': prediction_result.get('top_5_predictions', [{}])[3].get('virus', '') if len(prediction_result.get('top_5_predictions', [])) > 3 else '', 'top_4_confidence': prediction_result.get('top_5_predictions', [{}])[3].get('confidence', 0) if len(prediction_result.get('top_5_predictions', [])) > 3 else 0, 'top_5_virus': prediction_result.get('top_5_predictions', [{}])[4].get('virus', '') if len(prediction_result.get('top_5_predictions', [])) > 4 else '', 'top_5_confidence': prediction_result.get('top_5_predictions', [{}])[4].get('confidence', 0) if len(prediction_result.get('top_5_predictions', [])) > 4 else 0 } # Prepare complete document for single collection document = { # Patient information (flat structure) **readable_patient_info, # Symptoms (flat structure - each symptom as separate field) **symptoms_readable, # Predictions (flat structure) **prediction_readable, # Validation fields (empty initially, filled when validated) 'validation_status': 'pending', 'actual_virus_name': '', 'actual_virus_category': '', 'validation_confidence_level': '', 'validation_notes': '', 'validated_at': None, 'validated_by': '', # Metadata 'prediction_timestamp': datetime.utcnow(), 'model_primary': models_info.get('model1', '') if models_info else '', 'model_secondary': models_info.get('model2', '') if models_info else '', 'app_version': '2.0' } # Insert document result = collection.insert_one(document) # Update usage statistics self._update_usage_stats() logger.info(f"Prediction saved with Patient ID: {patient_id}, Document ID: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"Error saving prediction: {e}") return None def _get_month_name(self, month_num: int) -> str: """Convert month number to month name""" months = ['', 'January', 'February', 'March', 'April', 'May', 'June', 'July', 'August', 'September', 'October', 'November', 'December'] return months[month_num] if 1 <= month_num <= 12 else 'Unknown' def _transform_symptoms_to_readable(self, patient_data: Dict) -> Dict: """Transform symptom flags to human-readable flat structure for CSV export""" # Define all possible symptoms (from SYMPTOM_GROUPS in app.py) all_symptoms = [ 'HEADACHE', 'IRRITABILITY', 'ALTERED SENSORIUM', 'SOMNOLENCE', 'NECK RIGIDITY', 'SEIZURES', 'DIARRHEA', 'DYSENTERY', 'NAUSEA', 'VOMITING', 'ABDOMINAL PAIN', 'MALAISE', 'MYALGIA', 'ARTHRALGIA', 'CHILLS', 'RIGORS', 'FEVER', 'BREATHLESSNESS', 'COUGH', 'RHINORRHEA', 'SORE THROAT', 'BULLAE', 'PAPULAR RASH', 'PUSTULAR RASH', 'MUSCULAR RASH', 'MACULOPAPULAR RASH', 'ESCHAR', 'DARK URINE', 'HEPATOMEGALY', 'JAUNDICE', 'RED EYE', 'DISCHARGE EYES', 'CRUSHING EYES' ] # Create flat symptom structure symptoms_dict = {} for symptom in all_symptoms: readable_name = f"symptom_{symptom.lower().replace(' ', '_')}" symptoms_dict[readable_name] = 'Yes' if patient_data.get(symptom, 0) == 1 else 'No' return symptoms_dict def save_patient(self, patient_data: Dict) -> Optional[str]: """ Save patient information to database Args: patient_data: Patient demographic and clinical information Returns: Document ID if successful, None otherwise """ try: if self.db is None: logger.error("Database not initialized") return None collection = self.db['patients'] # Add metadata document = { **patient_data, 'created_at': datetime.utcnow(), 'updated_at': datetime.utcnow() } result = collection.insert_one(document) logger.info(f"Patient saved with ID: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"Error saving patient: {e}") return None def get_prediction_history(self, limit: int = 100, patient_id: str = None) -> List[Dict]: """ Retrieve prediction history Args: limit: Maximum number of records to return patient_id: Filter by specific patient ID Returns: List of prediction records """ try: if self.db is None: return [] collection = self.db['predictions'] # Build query query = {} if patient_id: query['patient_data.patient_id'] = patient_id # Get records cursor = collection.find(query).sort('timestamp', -1).limit(limit) records = list(cursor) # Convert ObjectId to string for JSON serialization for record in records: record['_id'] = str(record['_id']) return records except Exception as e: logger.error(f"Error retrieving prediction history: {e}") return [] def get_usage_statistics(self) -> Dict: """ Get usage statistics Returns: Dictionary with usage statistics """ try: if self.db is None: return {} predictions_collection = self.db['predictions'] # Get total predictions total_predictions = predictions_collection.count_documents({}) # Get predictions by virus type pipeline = [ { '$group': { '_id': '$prediction_result.predicted_virus', 'count': {'$sum': 1} } }, {'$sort': {'count': -1}} ] virus_stats = list(predictions_collection.aggregate(pipeline)) # Get predictions by date (last 30 days) from datetime import timedelta thirty_days_ago = datetime.utcnow() - timedelta(days=30) daily_pipeline = [ { '$match': { 'timestamp': {'$gte': thirty_days_ago} } }, { '$group': { '_id': { '$dateToString': { 'format': '%Y-%m-%d', 'date': '$timestamp' } }, 'count': {'$sum': 1} } }, {'$sort': {'_id': 1}} ] daily_stats = list(predictions_collection.aggregate(daily_pipeline)) return { 'total_predictions': total_predictions, 'virus_distribution': virus_stats, 'daily_predictions': daily_stats, 'last_updated': datetime.utcnow().isoformat() } except Exception as e: logger.error(f"Error getting usage statistics: {e}") return {} def _update_usage_stats(self): """Update daily usage statistics""" try: if self.db is None: return collection = self.db['usage_stats'] today = datetime.utcnow().date().isoformat() # Update or create today's stats collection.update_one( {'date': today}, { '$inc': {'prediction_count': 1}, '$set': {'last_updated': datetime.utcnow()} }, upsert=True ) except Exception as e: logger.error(f"Error updating usage stats: {e}") def save_validation(self, validation_data: Dict) -> Optional[str]: """ Save medical validation data within the same collection document Args: validation_data: Validation information including actual diagnosis Returns: Document ID if successful, None otherwise """ try: if self.db is None: logger.error("Database not initialized") return None # Use the single collection collection = self.db['virus_predictions'] prediction_id = validation_data.get('prediction_id') if not prediction_id: logger.error("No prediction_id provided in validation data") return None # Convert string ID to ObjectId for MongoDB from bson import ObjectId try: object_id = ObjectId(prediction_id) except Exception as e: logger.error(f"Invalid prediction_id format: {e}") return None # Prepare validation fields for flat structure validation_fields = { 'validation_status': 'validated', 'actual_virus_name': validation_data.get('actual_virus_name', ''), 'actual_virus_category': 'Main' if validation_data.get('actual_virus_key', '').startswith('main_') else 'Other', 'validation_confidence_level': validation_data.get('confidence_level', ''), 'validation_notes': validation_data.get('notes', ''), 'validated_at': datetime.utcnow(), 'validated_by': validation_data.get('validated_by', 'Medical Professional'), 'validation_accuracy': 'Correct' if validation_data.get('is_correct', False) else 'Incorrect' } # Update the document with validation data result = collection.update_one( {'_id': object_id}, { '$set': validation_fields, '$currentDate': {'last_updated': True} } ) if result.modified_count > 0: logger.info(f"Validation added to prediction ID: {prediction_id}") return prediction_id else: logger.warning(f"No prediction found with ID: {prediction_id}") return None except Exception as e: logger.error(f"Error saving validation: {e}") return None def get_validation_stats(self) -> Dict: """ Get validation statistics for data collection and analysis Note: This is for research/improvement purposes, not system accuracy calculation Returns: Dictionary containing validation collection statistics """ try: if self.db is None: return {'status': 'error', 'message': 'Database not initialized'} collection = self.db['predictions'] # Count total predictions with validation data total_validations = collection.count_documents({'validation.validated': True}) # Get validation distribution by actual virus pipeline = [ { '$match': {'validation.validated': True} }, { '$group': { '_id': '$validation.actual_virus_name', 'count': {'$sum': 1} } }, {'$sort': {'count': -1}} ] validation_distribution = list(collection.aggregate(pipeline)) # Get validation confidence distribution confidence_pipeline = [ { '$match': {'validation.validated': True} }, { '$group': { '_id': '$validation.confidence_level', 'count': {'$sum': 1} } } ] confidence_stats = list(collection.aggregate(confidence_pipeline)) return { 'status': 'success', 'total_validations': total_validations, 'validation_distribution': validation_distribution, 'confidence_distribution': confidence_stats, 'last_updated': datetime.utcnow().isoformat(), 'note': 'This data is for research and model improvement purposes' } except Exception as e: logger.error(f"Error getting validation stats: {e}") return {'status': 'error', 'message': str(e)} def export_to_csv(self, limit: int = None) -> Optional[pd.DataFrame]: """ Export all data to CSV-ready DataFrame format Args: limit: Maximum number of records to export (None for all) Returns: pandas DataFrame ready for CSV export """ try: if self.db is None: logger.error("Database not initialized") return None collection = self.db['virus_predictions'] # Build query to exclude MongoDB internal fields projection = {'_id': 0, 'encoded_data': 0} # Exclude internal fields # Get records if limit: cursor = collection.find({}, projection).sort('prediction_timestamp', -1).limit(limit) else: cursor = collection.find({}, projection).sort('prediction_timestamp', -1) records = list(cursor) if not records: logger.warning("No records found for export") return pd.DataFrame() # Convert to DataFrame df = pd.DataFrame(records) # Format timestamps for better readability if 'prediction_timestamp' in df.columns: df['prediction_timestamp'] = df['prediction_timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S') if 'validated_at' in df.columns: df['validated_at'] = df['validated_at'].apply( lambda x: x.strftime('%Y-%m-%d %H:%M:%S') if pd.notnull(x) else '' ) # Reorder columns for better CSV structure column_order = [ 'patient_id', 'age', 'sex', 'patient_type', 'state_name', 'district_name', 'syndrome_name', 'syndrome_specification', 'duration_of_illness_days', 'month_name', 'year', 'prediction_timestamp' ] # Add symptom columns symptom_cols = [col for col in df.columns if col.startswith('symptom_')] column_order.extend(sorted(symptom_cols)) # Add prediction columns prediction_cols = [ 'predicted_virus_name', 'prediction_confidence_percent', 'top_1_virus', 'top_1_confidence', 'top_2_virus', 'top_2_confidence', 'top_3_virus', 'top_3_confidence', 'top_4_virus', 'top_4_confidence', 'top_5_virus', 'top_5_confidence' ] column_order.extend(prediction_cols) # Add validation columns validation_cols = [ 'validation_status', 'actual_virus_name', 'actual_virus_category', 'validation_confidence_level', 'validation_notes', 'validated_at', 'validated_by', 'validation_accuracy' ] column_order.extend(validation_cols) # Add metadata columns metadata_cols = ['model_primary', 'model_secondary', 'app_version'] column_order.extend(metadata_cols) # Reorder DataFrame columns existing_cols = [col for col in column_order if col in df.columns] remaining_cols = [col for col in df.columns if col not in existing_cols] final_column_order = existing_cols + remaining_cols df = df[final_column_order] logger.info(f"Exported {len(df)} records to DataFrame") return df except Exception as e: logger.error(f"Error exporting to CSV: {e}") return None def health_check(self) -> Dict: """ Perform health check on database connection and operations Returns: Health check results """ try: # Check if we have a database instance (connection already established) if self.db is None: return { 'status': 'error', 'message': 'Database instance not available', 'details': { 'timestamp': datetime.utcnow().isoformat() } } # Test basic operations using existing connection try: # Simple test - count documents in virus_predictions collection predictions_count = self.db['virus_predictions'].count_documents({}) return { 'status': 'healthy', 'message': 'All database operations working', 'details': { 'connection': 'OK', 'total_predictions': predictions_count, 'timestamp': datetime.utcnow().isoformat() } } except Exception as op_error: return { 'status': 'error', 'message': f'Database operations failed: {str(op_error)}', 'details': { 'error': str(op_error), 'timestamp': datetime.utcnow().isoformat() } } except Exception as e: return { 'status': 'error', 'message': f'Health check failed: {str(e)}', 'details': { 'error': str(e), 'timestamp': datetime.utcnow().isoformat() } } # Global data handler instance data_handler = DataHandler() # Convenience functions for use in app.py def save_prediction_to_db(patient_data: Dict, prediction_result: Dict, models_info: Dict = None, state_name: str = None, district_name: str = None) -> Optional[str]: """Save prediction to database""" return data_handler.save_prediction(patient_data, prediction_result, models_info, state_name, district_name) def save_validation_to_db(validation_data: Dict) -> Optional[str]: """Save validation to database""" return data_handler.save_validation(validation_data) def get_db_health() -> Dict: """Get database health status""" return data_handler.health_check() def get_prediction_stats() -> Dict: """Get prediction usage statistics""" return data_handler.get_usage_statistics() def get_validation_stats() -> Dict: """Get validation statistics""" return data_handler.get_validation_stats() def export_data_to_csv(limit: int = None) -> Optional[pd.DataFrame]: """Export data to CSV-ready DataFrame""" return data_handler.export_to_csv(limit)