# src/metadata_template_generator/utils/bruker_metadata_extractor.py
"""
Extract and map Bruker OPUS parameters to template metadata format.
"""
from typing import Dict, Any, Optional, List
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

class BrukerMetadataExtractor:
    """Extract metadata from Bruker OPUS data and map to template format."""
    
    # Updated mapping with LOWERCASE keys (as they appear in the Parameters object)
    BRUKER_TO_TEMPLATE_MAPPING = {
        # Instrument Configuration
        'ins': 'spectrometer_model',        # Instrument Type
        'srn': 'spectrometer_serial_number', # Instrument Serial Number
        
        # Optical Parameters  
        'dtc': 'detector_type',    # Detector Setting
        'src': 'source_type',      # Source Setting
        'bms': 'beamsplitter',     # Beamsplitter Setting
        'apt': 'aperture_setting', # Aperture Setting
        'opf': 'optical_filter',   # Optical Filter Setting
        'pgn': 'preamplifier_gain',# Preamplifier Gain
        'vel': 'scanner_velocity', # Scanner Velocity
        
        # Acquisition Parameters
        'nss': 'number_of_scans',     # Sample Scans
        'nsr': 'background_scans',    # Background Scans (from reference)
        'ars': 'background_scans',    # Alternative: Number of Background Scans
        'res': 'resolution',          # Resolution
        'lfw': 'spectral_range_start', # Wanted Low Frequency Limit
        'hfw': 'spectral_range_end',   # Wanted High Frequency Limit
        'dur': 'measurement_duration', # Scan time (sec)
        
        # Fourier Transform Parameters
        'apf': 'apodization_function',   # Apodization Function
        'zff': 'zero_filling_factor',   # Zero Filling Factor
        'phz': 'phase_correction_mode', # Phase Correction Mode
        'phr': 'phase_resolution',      # Phase Resolution
        
        # Sample Information
        'snm': 'sample_name',    # Sample Name
        'sfm': 'sample_form',    # Sample Form
        'cnm': 'operator_name',  # Operator Name
        'exp': 'experiment_name', # Experiment
        
        # Environmental/Quality
        'tsc': 'scanner_temperature',  # Scanner Temperature
        'hum': 'humidity_level',       # Relative Humidity
        'srt': 'measurement_timestamp', # Start time (sec)
        'uid': 'measurement_uuid',      # Universally Unique Identifier
        
        # Additional useful parameters
        'acc': 'accessory',           # Accessory (for ATR detection)
        'chn': 'measurement_channel', # Measurement Channel
    }
    
    # Value transformations for specific fields
    VALUE_TRANSFORMS = {
        'apodization_function': {
            'B3': 'Blackman-Harris 3-term',
            'BH': 'Blackman-Harris 4-term', 
            'HG': 'Happ-Genzel',
            'TR': 'Triangular',
            'BX': 'Boxcar',
        },
        'phase_correction_mode': {
            'ML': 'Mertz-Lynch',
            'PK': 'Peak picking',
            'NO': 'None',
        },
    }
    
    def extract_parameters_from_opus(self, opus_data) -> Dict[str, Any]:
        """
        Extract parameters from opus_data object handling Parameters objects.
        
        Args:
            opus_data: Loaded brukeropus data object
            
        Returns:
            Dictionary of all found parameters
        """
        all_params = {}
        
        # Method 1: Handle Parameters objects using .items()
        if hasattr(opus_data, 'params') and opus_data.params is not None:
            try:
                sample_params = dict(opus_data.params.items())
                all_params.update(sample_params)
                logger.info(f"Found {len(sample_params)} sample parameters")
            except Exception as e:
                logger.warning(f"Could not extract sample parameters: {e}")
            
        if hasattr(opus_data, 'rf_params') and opus_data.rf_params is not None:
            try:
                ref_params = dict(opus_data.rf_params.items())
                # Reference params first, then sample params override
                temp_params = ref_params.copy()
                temp_params.update(all_params)
                all_params = temp_params
                logger.info(f"Found {len(ref_params)} reference parameters")
            except Exception as e:
                logger.warning(f"Could not extract reference parameters: {e}")
        
        logger.info(f"Total extracted parameters: {len(all_params)}")
        
        if not all_params:
            logger.error("❌ No parameters could be extracted from OPUS file")
        else:
            # Debug: Show what we found
            logger.info("Key parameters found:")
            for key in ['ins', 'srn', 'nss', 'res', 'snm', 'dtc', 'src']:
                if key in all_params:
                    logger.info(f"  {key}: {all_params[key]}")
        
        return all_params
    
    def extract_metadata(
        self, 
        opus_data, 
        template_field: List[Dict[str, Any]],
        additional_metadata: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Extract metadata from Bruker OPUS data object."""
        
        logger.info("Extracting Bruker metadata...")
        
        # Extract all parameters
        all_params = self.extract_parameters_from_opus(opus_data)
        
        if not all_params:
            logger.error("No parameters could be extracted from OPUS data")
            return additional_metadata or {}
        
        extracted = {}
        
        # Map Bruker parameters to template fields
        for bruker_key, template_field in self.BRUKER_TO_TEMPLATE_MAPPING.items():
            if bruker_key in all_params:
                raw_value = all_params[bruker_key]
                
                # Clean up the value
                if isinstance(raw_value, str):
                    raw_value = raw_value.strip()
                    if not raw_value:  # Empty string
                        continue
                
                # Apply value transformations
                if template_field in self.VALUE_TRANSFORMS:
                    transform_map = self.VALUE_TRANSFORMS[template_field]
                    transformed_value = transform_map.get(str(raw_value), raw_value)
                    extracted[template_field] = transformed_value
                else:
                    extracted[template_field] = raw_value
                
                logger.debug(f"Mapped {bruker_key} -> {template_field}: {raw_value}")
        
        # Extract spectral range from data
        extracted = self._extract_spectral_range(opus_data, extracted)
        
        # Determine measurement technique
        extracted = self._infer_measurement_technique(all_params, extracted)
        
        # Add calculated fields
        extracted = self._add_calculated_fields(all_params, extracted)
        
        # Add user-provided metadata
        if additional_metadata:
            extracted.update(additional_metadata)
        
        logger.info(f"Successfully extracted {len(extracted)} metadata fields")
        return extracted
    
    def _extract_spectral_range(self, opus_data, metadata: Dict[str, Any]) -> Dict[str, Any]:
        """Extract spectral range from spectrum data."""
        try:
            # Try different ways to get wavenumber data
            wavenumbers = None
            
            # Method 1: Check if there's spectral data
            if hasattr(opus_data, 'a') and opus_data.a is not None:
                # 'a' is typically the absorbance spectrum
                spectrum = opus_data.a
                if hasattr(spectrum, 'x'):
                    wavenumbers = spectrum.x
                elif hasattr(spectrum, 'wavenumbers'):
                    wavenumbers = spectrum.wavenumbers
            
            # Method 2: Try other spectrum types
            if wavenumbers is None:
                for spectrum_attr in ['sm', 'igrf', 'phsm']:  # Other spectrum types
                    if hasattr(opus_data, spectrum_attr):
                        spectrum = getattr(opus_data, spectrum_attr)
                        if spectrum is not None and hasattr(spectrum, 'x'):
                            wavenumbers = spectrum.x
                            break
            
            if wavenumbers is not None:
                import numpy as np
                wavenumbers = np.asarray(wavenumbers)
                metadata['actual_spectral_range_start'] = float(wavenumbers.min())
                metadata['actual_spectral_range_end'] = float(wavenumbers.max())
                metadata['actual_data_points'] = len(wavenumbers)
                logger.info(f"Extracted spectral range: {wavenumbers.min():.1f} - {wavenumbers.max():.1f} cm⁻¹")
            else:
                logger.warning("Could not extract wavenumber data from spectrum")
                
        except Exception as e:
            logger.warning(f"Error extracting spectral range: {e}")
        
        return metadata
    
    def _infer_measurement_technique(
        self, 
        params: Dict[str, Any], 
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Infer measurement technique from parameters."""
        
        # Check for ATR (accessory parameter)
        acc = params.get('acc', '').upper()
        if 'ATR' in acc or 'GOLDEN GATE' in acc:
            metadata['measurement_technique'] = 'ATR-FTIR'
        # Check for reflection
        elif 'REFLECTION' in acc or 'REFL' in acc:
            metadata['measurement_technique'] = 'Reflection FTIR'
        # Check measurement channel for external detectors (might indicate RAIRS)
        elif 'external' in params.get('dtc', '').lower():
            metadata['measurement_technique'] = 'External Reflection FTIR'
        # Default to transmission
        else:
            metadata['measurement_technique'] = 'Transmission FTIR'
        
        return metadata
    
    def _add_calculated_fields(
        self, 
        params: Dict[str, Any], 
        metadata: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Add calculated/derived fields."""
        
        # Calculate total measurement time
        try:
            scan_time = float(params.get('dur', 0))
            num_scans = int(params.get('nss', 1) or 1)
            bg_scans = int(params.get('nsr', params.get('ars', 0)) or 0)
            
            if scan_time and num_scans:
                total_time = scan_time * (num_scans + bg_scans) / num_scans
                metadata['total_measurement_time'] = round(total_time, 2)
        except (ValueError, TypeError):
            pass
        
        # Convert numeric string values to proper numbers
        for key, value in metadata.items():
            if isinstance(value, str) and value.replace('.', '').replace('-', '').isdigit():
                try:
                    if '.' in value:
                        metadata[key] = float(value)
                    else:
                        metadata[key] = int(value)
                except ValueError:
                    pass
        
        return metadata

# Updated convenience function (same as before)
def extract_bruker_metadata_for_curve_fitting(
    opus_data,
    template_result: Dict[str, Any],
    user_metadata: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """Extract Bruker metadata optimized for curve fitting analysis."""
    
    extractor = BrukerMetadataExtractor()
    
    # Extract from Bruker data
    template_fields = template_result.get('fields', [])
    extracted = extractor.extract_metadata(opus_data, template_fields, user_metadata)
    
    # Create curve fitting specific structure
    curve_metadata = {
        "instrument": f"{extracted.get('spectrometer_model', 'Unknown')} S/N:{extracted.get('spectrometer_serial_number', 'Unknown')}",
        "sample": extracted.get('sample_name', extracted.get('sample', 'Unknown')),
        "sample_form": extracted.get('sample_form', ''),
        "technique": extracted.get('measurement_technique', extracted.get('technique', 'FTIR')),
        "detector": extracted.get('detector_type', ''),
        "source": extracted.get('source_type', ''),
        "spectral_range_cm1": [
            extracted.get('actual_spectral_range_start', extracted.get('spectral_range_start')),
            extracted.get('actual_spectral_range_end', extracted.get('spectral_range_end'))
        ],
        "resolution_cm1": extracted.get('resolution'),
        "scans": {
            "sample": extracted.get('number_of_scans'),
            "background": extracted.get('background_scans'),
        },
        "optics": {
            "beamsplitter": extracted.get('beamsplitter'),
            "aperture": extracted.get('aperture_setting'),
            "detector": extracted.get('detector_type'),
            "source": extracted.get('source_type')
        },
        "processing": {
            "apodization": extracted.get('apodization_function'),
            "zero_filling": extracted.get('zero_filling_factor'),
            "phase_correction": extracted.get('phase_correction_mode')
        },
        "quality": {
            "scanner_temperature": extracted.get('scanner_temperature'),
            "humidity": extracted.get('humidity_level'),
            "measurement_duration": extracted.get('total_measurement_time')
        },
        # Include all user metadata
        **{k: v for k, v in (user_metadata or {}).items() if k not in ['technique']}
    }
    
    # Remove None values but keep empty containers
    return _remove_none_values(curve_metadata)

def _remove_none_values(d):
    """Recursively remove None values."""
    if isinstance(d, dict):
        cleaned = {}
        for k, v in d.items():
            cleaned_v = _remove_none_values(v)
            if cleaned_v is not None and cleaned_v != "" and cleaned_v != [None, None]:
                cleaned[k] = cleaned_v
        return cleaned
    elif isinstance(d, list):
        cleaned = [_remove_none_values(x) for x in d if _remove_none_values(x) is not None]
        return cleaned if any(x is not None for x in cleaned) else []
    else:
        return d
