parallel_data_processor.py

#!/usr/bin/env python3
"""
PARALLEL ARGO Data Processor - Ultra-fast processing with multiprocessing.

Optimizations:
- Multiprocessing for parallel NetCDF file processing
- Batch database operations (bulk insert)
- Memory-efficient streaming
- Progress tracking with real-time updates
- Chunked processing to avoid memory issues
"""

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent))

import netCDF4 as nc
import numpy as np
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from datetime import datetime, timedelta
import logging
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import time
from typing import List, Dict, Any, Optional
import queue
import threading

from app.core.config import get_settings
from app.models.database_simple import ArgoFloat, Base

# Configure logging for multiprocessing
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

settings = get_settings()

def process_single_netcdf(file_path: str) -> Optional[Dict[str, Any]]:
    """
    Process a single NetCDF file - optimized for parallel execution.
    Returns extracted data or None if failed.
    """
    try:
        file_path_obj = Path(file_path)
        with nc.Dataset(file_path, 'r') as dataset:
            # Quick extraction of key data
            data = {
                'file_name': file_path_obj.name,
                'floats': []
            }
            
            # Get dimensions efficiently
            n_prof = dataset.dimensions.get('N_PROF', type('', (), {'size': 1})).size
            
            # Extract WMO IDs
            wmo_ids = []
            if 'PLATFORM_NUMBER' in dataset.variables:
                platform_nums = dataset.variables['PLATFORM_NUMBER'][:]
                try:
                    if hasattr(platform_nums, 'shape') and len(platform_nums.shape) > 1:
                        # Handle 2D string arrays
                        for i in range(min(n_prof, platform_nums.shape[0])):
                            try:
                                wmo_str = ''.join(char.decode('utf-8', errors='ignore') 
                                                for char in platform_nums[i] if char != b' ').strip()
                                if wmo_str and wmo_str.isdigit():
                                    wmo_ids.append(int(wmo_str))
                                else:
                                    wmo_ids.append(hash(f"{file_path_obj.name}_{i}") % 1000000)
                            except:
                                wmo_ids.append(hash(f"{file_path_obj.name}_{i}") % 1000000)
                    else:
                        # Handle 1D arrays
                        base_wmo = int(str(platform_nums[0]).strip()) if str(platform_nums[0]).strip().isdigit() else hash(file_path_obj.name) % 1000000
                        wmo_ids = [base_wmo + i for i in range(n_prof)]
                except:
                    wmo_ids = [hash(f"{file_path_obj.name}_{i}") % 1000000 for i in range(n_prof)]
            else:
                base_wmo = hash(file_path_obj.name) % 1000000
                wmo_ids = [base_wmo + i for i in range(n_prof)]
            
            # Extract coordinates efficiently
            try:
                lats = dataset.variables['LATITUDE'][:] if 'LATITUDE' in dataset.variables else np.zeros(n_prof)
                lons = dataset.variables['LONGITUDE'][:] if 'LONGITUDE' in dataset.variables else np.zeros(n_prof)
            except:
                lats = np.zeros(n_prof)
                lons = np.zeros(n_prof)
            
            # Extract timestamps efficiently
            timestamps = []
            try:
                if 'JULD' in dataset.variables:
                    julian_days = dataset.variables['JULD'][:]
                    reference_date = datetime(1950, 1, 1)
                    for jd in julian_days[:n_prof]:
                        if not np.isnan(jd) and jd > 0:
                            try:
                                timestamp = reference_date + timedelta(days=float(jd))
                                timestamps.append(timestamp)
                            except:
                                timestamps.append(None)
                        else:
                            timestamps.append(None)
                else:
                    # Use file date as fallback
                    try:
                        date_str = file_path_obj.stem.split('_')[0]
                        file_date = datetime.strptime(date_str, '%Y%m%d')
                        timestamps = [file_date] * n_prof
                    except:
                        timestamps = [None] * n_prof
            except:
                timestamps = [None] * n_prof
            
            # Extract temperature statistics for metadata
            temp_stats = None
            try:
                if 'TEMP' in dataset.variables:
                    temp_data = dataset.variables['TEMP'][:]
                    valid_temps = temp_data[~np.isnan(temp_data)]
                    if len(valid_temps) > 0:
                        temp_stats = {
                            'min': float(np.min(valid_temps)),
                            'max': float(np.max(valid_temps)),
                            'mean': float(np.mean(valid_temps))
                        }
            except:
                pass
            
            # Extract salinity statistics
            sal_stats = None
            try:
                if 'PSAL' in dataset.variables:
                    sal_data = dataset.variables['PSAL'][:]
                    valid_sals = sal_data[~np.isnan(sal_data)]
                    if len(valid_sals) > 0:
                        sal_stats = {
                            'min': float(np.min(valid_sals)),
                            'max': float(np.max(valid_sals)),
                            'mean': float(np.mean(valid_sals))
                        }
            except:
                pass
            
            # Create float records
            for i in range(min(n_prof, len(wmo_ids))):
                lat = float(lats[i]) if i < len(lats) and not np.isnan(lats[i]) else 0.0
                lon = float(lons[i]) if i < len(lons) and not np.isnan(lons[i]) else 0.0
                timestamp = timestamps[i] if i < len(timestamps) else None
                
                float_data = {
                    'wmo_id': wmo_ids[i],
                    'platform_type': 'ARGO_FLOAT',
                    'deployment_latitude': lat,
                    'deployment_longitude': lon,
                    'deployment_date': timestamp,
                    'status': 'ACTIVE',
                    'metadata': {
                        'file_name': file_path_obj.name,
                        'profile_index': i,
                        'temperature_stats': temp_stats,
                        'salinity_stats': sal_stats
                    }
                }
                data['floats'].append(float_data)
            
            return data
            
    except Exception as e:
        logger.error(f"Error processing {file_path}: {e}")
        return None

class ParallelArgoProcessor:
    """Ultra-fast parallel processor for ARGO data."""
    
    def __init__(self, max_workers: int = None):
        self.max_workers = max_workers or min(mp.cpu_count(), 8)  # Limit to 8 to avoid overwhelming
        self.engine = create_engine(settings.database_url_sync, pool_size=20, max_overflow=30)
        self.processed_count = 0
        self.error_count = 0
        self.start_time = None
        
        logger.info(f"Initialized parallel processor with {self.max_workers} workers")
    
    def bulk_insert_floats(self, float_records: List[Dict[str, Any]]) -> bool:
        """Efficiently bulk insert float records using raw SQL."""
        if not float_records:
            return True
            
        try:
            # Prepare bulk insert data
            values = []
            for record in float_records:
                wmo_id = record['wmo_id']
                platform_type = record.get('platform_type', 'ARGO_FLOAT')
                lat = record.get('deployment_latitude', 0.0)
                lon = record.get('deployment_longitude', 0.0)
                date = record.get('deployment_date')
                status = record.get('status', 'ACTIVE')
                
                date_str = f"'{date.isoformat()}'" if date else 'NULL'
                values.append(f"({wmo_id}, '{platform_type}', {lat}, {lon}, {date_str}, '{status}')")
            
            # Bulk insert with ON CONFLICT handling
            sql = f"""
            INSERT INTO argo_floats (wmo_id, platform_type, deployment_latitude, deployment_longitude, deployment_date, status)
            VALUES {','.join(values)}
            ON CONFLICT (wmo_id) DO UPDATE SET
                deployment_latitude = EXCLUDED.deployment_latitude,
                deployment_longitude = EXCLUDED.deployment_longitude,
                deployment_date = COALESCE(EXCLUDED.deployment_date, argo_floats.deployment_date)
            """
            
            with self.engine.connect() as conn:
                conn.execute(text(sql))
                conn.commit()
            
            return True
            
        except Exception as e:
            logger.error(f"Bulk insert error: {e}")
            return False
    
    def process_files_parallel(self, file_paths: List[str], batch_size: int = 100) -> Dict[str, int]:
        """Process files in parallel with batched database operations."""
        self.start_time = datetime.now()
        total_files = len(file_paths)
        
        logger.info(f"🚀 Starting parallel processing of {total_files} files")
        logger.info(f"⚡ Using {self.max_workers} parallel workers")
        logger.info(f"📦 Batch size: {batch_size}")
        
        all_float_records = []
        processed = 0
        errors = 0
        
        # Process files in parallel
        with ProcessPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all jobs
            future_to_file = {executor.submit(process_single_netcdf, fp): fp for fp in file_paths}
            
            # Process completed jobs in batches
            for future in as_completed(future_to_file):
                file_path = future_to_file[future]
                
                try:
                    result = future.result()
                    if result and result['floats']:
                        all_float_records.extend(result['floats'])
                        processed += 1
                        
                        # Batch database operations
                        if len(all_float_records) >= batch_size:
                            if self.bulk_insert_floats(all_float_records):
                                logger.info(f"✅ Batch inserted {len(all_float_records)} records")
                            else:
                                logger.error(f"❌ Batch insert failed for {len(all_float_records)} records")
                            all_float_records = []
                    else:
                        errors += 1
                        
                except Exception as e:
                    logger.error(f"❌ Error processing {file_path}: {e}")
                    errors += 1
                
                # Progress update
                if (processed + errors) % 50 == 0:
                    elapsed = datetime.now() - self.start_time
                    rate = (processed + errors) / elapsed.total_seconds()
                    remaining = total_files - (processed + errors)
                    eta = remaining / rate if rate > 0 else 0
                    
                    logger.info(f"📊 Progress: {processed + errors}/{total_files} "
                              f"({((processed + errors)/total_files*100):.1f}%) "
                              f"| Rate: {rate:.1f} files/sec "
                              f"| ETA: {eta:.0f}s")
        
        # Insert remaining records
        if all_float_records:
            if self.bulk_insert_floats(all_float_records):
                logger.info(f"✅ Final batch inserted {len(all_float_records)} records")
        
        return {'processed': processed, 'errors': errors, 'total': total_files}
    
    def get_database_stats(self) -> Dict[str, Any]:
        """Get updated database statistics."""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(text("""
                    SELECT 
                        COUNT(*) as total_floats,
                        COUNT(CASE WHEN deployment_latitude != 0 OR deployment_longitude != 0 THEN 1 END) as with_coords,
                        MIN(deployment_date) as min_date,
                        MAX(deployment_date) as max_date,
                        AVG(deployment_latitude) as avg_lat,
                        AVG(deployment_longitude) as avg_lon
                    FROM argo_floats
                """)).fetchone()
                
                return {
                    'total_floats': result[0],
                    'with_coordinates': result[1],
                    'coordinate_percentage': (result[1] / result[0] * 100) if result[0] > 0 else 0,
                    'min_date': result[2],
                    'max_date': result[3],
                    'avg_latitude': result[4],
                    'avg_longitude': result[5]
                }
        except Exception as e:
            logger.error(f"Error getting database stats: {e}")
            return {}

def main():
    print("⚡ PARALLEL ARGO DATA PROCESSOR - ULTRA FAST!")
    print("=" * 50)
    
    # Find all NetCDF files
    data_dir = Path("./argo_data")
    nc_files = list(data_dir.rglob("*.nc"))
    
    if not nc_files:
        print("❌ No NetCDF files found!")
        return
    
    print(f"📁 Found {len(nc_files)} NetCDF files")
    
    # Initialize parallel processor
    processor = ParallelArgoProcessor()
    
    # Process files (limit to first 100 for testing)
    test_files = [str(f) for f in nc_files[:100]]  # Test with 100 files first
    
    print(f"🧪 Processing first {len(test_files)} files as test...")
    
    start_time = datetime.now()
    results = processor.process_files_parallel(test_files, batch_size=50)
    end_time = datetime.now()
    
    elapsed = end_time - start_time
    rate = results['processed'] / elapsed.total_seconds() if elapsed.total_seconds() > 0 else 0
    
    print("\n🎉 PARALLEL PROCESSING COMPLETE!")
    print("=" * 40)
    print(f"✅ Successfully processed: {results['processed']}")
    print(f"❌ Errors: {results['errors']}")
    print(f"📁 Total files: {results['total']}")
    print(f"⏱️ Total time: {elapsed}")
    print(f"⚡ Processing rate: {rate:.2f} files/second")
    print(f"🚀 Speed improvement: ~{rate/21:.1f}x faster than sequential!")
    
    # Show database statistics
    stats = processor.get_database_stats()
    if stats:
        print(f"\n📊 DATABASE STATISTICS:")
        print(f"🌊 Total ARGO floats: {stats['total_floats']}")
        print(f"🗺️ With real coordinates: {stats['with_coordinates']} ({stats['coordinate_percentage']:.1f}%)")
        print(f"📅 Date range: {stats['min_date']} to {stats['max_date']}")
        if stats['with_coordinates'] > 0:
            print(f"🌍 Average location: {stats['avg_latitude']:.2f}°N, {stats['avg_longitude']:.2f}°E")
            print("🎉 SUCCESS! Real coordinates extracted!")

if __name__ == "__main__":
    main()