data_registry.py

# data_registry.py
import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional, Union
import os
import json

class DataRegistry:
    def __init__(self):
        self.data = {}
        self.metadata = {}
        self.healthcare_metadata = {}
        self.derived_columns = {}  # Track derived columns per file
    
    def add_path(self, path: str) -> bool:
        """Add a data file to the registry with dynamic processing."""
        try:
            file_name = os.path.basename(path)
            file_ext = os.path.splitext(file_name)[1].lower()
            
            # Read file based on extension
            if file_ext == '.csv':
                df = pd.read_csv(path)
            elif file_ext in ['.xlsx', '.xls']:
                df = pd.read_excel(path)
            elif file_ext == '.json':
                with open(path, 'r') as f:
                    data = json.load(f)
                df = pd.json_normalize(data)
            elif file_ext in ['.parquet']:
                df = pd.read_parquet(path)
            else:
                print(f"Unsupported file type: {file_ext}")
                return False
            
            # Standardize column names
            df.columns = [col.strip().lower().replace(' ', '_').replace('-', '_').replace('.', '_') for col in df.columns]
            
            # Store original dataframe
            self.data[file_name] = df.copy()
            
            # Initialize derived columns tracking
            self.derived_columns[file_name] = set()
            
            # Process healthcare data dynamically
            self._process_healthcare_data(file_name, df)
            
            # Basic metadata
            self.metadata[file_name] = {
                'type': file_ext,
                'columns': list(df.columns),
                'shape': df.shape,
                'sample': df.head(3).to_dict('records')
            }
            
            # Healthcare-specific metadata extraction
            self._extract_healthcare_metadata(file_name, df)
            
            return True
        except Exception as e:
            print(f"Error adding {path}: {e}")
            return False
    
    def _process_healthcare_data(self, file_name: str, df: pd.DataFrame):
        """Dynamically process healthcare data based on available columns."""
        # Dynamic column pattern matching
        column_patterns = {
            'facility_name': ['facility', 'name', 'hospital', 'site', 'location'],
            'facility_type': ['type', 'category', 'class', 'facility_type', 'odhf_facility_type'],
            'beds_current': ['current', '2023', '2024', 'beds_current', 'staffed_beds', 'capacity'],
            'beds_prev': ['prev', 'previous', '2022', 'beds_prev', 'previous_beds'],
            'zone': ['zone', 'region', 'area', 'district'],
            'province': ['province', 'state', 'territory'],
            'city': ['city', 'municipality', 'town'],
            'teaching_status': ['teaching', 'status', 'type', 'hospital_type']
        }
        
        # Map actual columns to standard names
        column_map = {}
        for standard_col, patterns in column_patterns.items():
            for col in df.columns:
                if any(pattern in col for pattern in patterns):
                    column_map[standard_col] = col
                    break
        
        # Create derived columns if we have the necessary base columns
        if 'beds_current' in column_map and 'beds_prev' in column_map:
            current_col = column_map['beds_current']
            prev_col = column_map['beds_prev']
            
            # Calculate bed change
            df['bed_change'] = df[current_col] - df[prev_col]
            self.derived_columns[file_name].add('bed_change')
            
            # Calculate percentage change (avoid division by zero)
            df['percent_change'] = df.apply(
                lambda row: (row['bed_change'] / row[prev_col] * 100) if row[prev_col] != 0 else 0, 
                axis=1
            )
            self.derived_columns[file_name].add('percent_change')
        
        # If we have facility_type but not in standard form, map it
        if 'facility_type' in column_map and column_map['facility_type'] != 'facility_type':
            df['facility_type'] = df[column_map['facility_type']]
            self.derived_columns[file_name].add('facility_type')
    
    def _extract_healthcare_metadata(self, file_name: str, df: pd.DataFrame):
        """Extract healthcare-specific metadata dynamically."""
        healthcare_meta = {}
        
        # Detect data type based on columns
        facility_cols = [col for col in df.columns if any(pattern in col for pattern in ['facility', 'name', 'site'])]
        bed_cols = [col for col in df.columns if any(pattern in col for pattern in ['bed', 'capacity'])]
        
        if facility_cols:
            healthcare_meta['data_type'] = 'facility_data'
            if 'facility_type' in df.columns:
                healthcare_meta['facility_types'] = df['facility_type'].value_counts().to_dict()
            if 'city' in df.columns:
                healthcare_meta['cities'] = df['city'].value_counts().head(10).to_dict()
        
        if bed_cols:
            healthcare_meta['data_type'] = 'bed_data'
            if 'zone' in df.columns:
                healthcare_meta['zones'] = df['zone'].unique().tolist()
            if 'teaching_status' in df.columns:
                healthcare_meta['teaching_status_counts'] = df['teaching_status'].value_counts().to_dict()
            
            # Check for derived metrics
            if 'bed_change' in df.columns:
                healthcare_meta['has_derived_metrics'] = True
        
        if healthcare_meta:
            self.healthcare_metadata[file_name] = healthcare_meta
    
    def get_derived_columns(self, file_name: str) -> set:
        """Get derived columns for a file."""
        return self.derived_columns.get(file_name, set())
    
    def find_column(self, file_name: str, patterns: List[str]) -> Optional[str]:
        """Find a column matching any of the given patterns."""
        df = self.get(file_name)
        if df is None:
            return None
        
        for col in df.columns:
            if any(pattern.lower() in col.lower() for pattern in patterns):
                return col
        return None
    
    def get_data_by_type(self, data_type: str) -> List[str]:
        """Get all files of a specific data type."""
        return [
            file_name for file_name, meta in self.healthcare_metadata.items()
            if meta.get('data_type') == data_type
        ]
    
    def names(self):
        return list(self.data.keys())
    
    def get(self, name):
        return self.data.get(name)
    
    def summarize_for_prompt(self) -> str:
        """Generate a summary of all data for prompt inclusion."""
        if not self.data:
            return "No data files registered."
        
        summary_parts = []
        for file_name in self.names():
            meta = self.metadata.get(file_name, {})
            health_meta = self.get_healthcare_metadata(file_name)
            
            summary_parts.append(f"File: {file_name}")
            summary_parts.append(f"Type: {meta.get('type', 'unknown')}")
            summary_parts.append(f"Columns: {', '.join(meta.get('columns', []))}")
            summary_parts.append(f"Shape: {meta.get('shape', 'unknown')}")
            
            if health_meta:
                summary_parts.append("Healthcare Context:")
                for key, value in health_meta.items():
                    if key != 'privacy_warning':
                        summary_parts.append(f"  {key}: {value}")
            
            summary_parts.append("")
        
        return "\n".join(summary_parts)
    
    def get_healthcare_metadata(self, name: str) -> Dict[str, Any]:
        """Get healthcare-specific metadata for a file."""
        return self.healthcare_metadata.get(name, {})
    
    def clear(self):
        self.data.clear()
        self.metadata.clear()
        self.healthcare_metadata.clear()
        self.derived_columns.clear()