import numpy as np
from pathlib import Path

def load_npz_files(data_dir, api_client=None, file_info_map=None, label_to_id=None):
    """
    Load .npz files from directory and combine into arrays.
    Khi file_info_map được truyền (danh sách file từ split), chỉ load các file có trong map
    để đảm bảo chỉ dùng dữ liệu đã filter theo split, không load thừa file .npz trong thư mục.
    
    Args:
        data_dir: Directory containing .npz files
        api_client: Optional KWSClient to get keyword info from sample_id
        file_info_map: Optional list of file info dicts (from split list_files). If provided,
                       only .npz files whose name is in this list are loaded (filter theo split).
        label_to_id: Optional existing label mapping dict. If provided, will reuse it
                     and add new labels if encountered. If None, creates new mapping.
    
    Returns:
        X: Features array
        y: Labels array
        label_to_id: Mapping from label to integer ID (reused or newly created)
    """
    all_npz = list(Path(data_dir).glob("*.npz"))
    
    # Build filename to metadata mapping if file_info_map provided
    filename_to_metadata = {}
    if file_info_map:
        for file_info in file_info_map:
            filename = file_info.get('file_name', '')
            metadata = {
                'sample_id': file_info.get('sample_id'),
                'derivative_label': file_info.get('derivative_label'),  # Use derivative_label from split
                'file_id': file_info.get('file_id'),
            }
            if filename:
                filename_to_metadata[filename] = metadata
        # Chỉ load file có trong split (đã filter) — tránh kéo hết .npz trong thư mục
        files = [f for f in all_npz if f.name in filename_to_metadata]
        if len(files) != len(all_npz):
            print(f"Loading only {len(files)} .npz files from split (ignored {len(all_npz) - len(files)} extra in dir)")
    else:
        files = all_npz
    print(f"Found {len(files)} .npz files")
    
    all_features = []
    all_labels = []
    
    # Reuse existing label mapping or create new one
    if label_to_id is None:
        label_to_id = {}
        current_id = 0
    else:
        # Use existing mapping, find max ID to continue from
        if label_to_id:
            current_id = max(label_to_id.values()) + 1
        else:
            current_id = 0
        print(f"Reusing existing label mapping with {len(label_to_id)} classes, continuing from ID {current_id}")
    
    # Cache for keyword lookups (fallback if derivative_label not available)
    keyword_cache = {}
    
    for file_path in files:
        try:
            data = np.load(file_path, allow_pickle=True)
            
            # Try common keys for features
            features = None
            if 'features' in data:
                features = data['features']
            elif 'feature' in data:
                features = data['feature']
            elif 'data' in data:
                features = data['data']
            elif 'X' in data:
                features = data['X']
            else:
                # Get first array that's not a label/metadata
                keys = list(data.keys())
                feature_key = None
                for key in keys:
                    if key not in ['label', 'labels', 'y', 'keyword', 'keyword_id', 'keyword_name', 
                                   'sample_id', 'feature_type_id', 'metadata']:
                        arr = data[key]
                        if isinstance(arr, np.ndarray) and arr.size > 0:
                            feature_key = key
                            break
                if feature_key:
                    features = data[feature_key]
            
            if features is None:
                print(f"Warning: Could not find features in {file_path}")
                continue
            
            # Get label - prioritize derivative_label from file_info_map
            label = None
            
            # Priority 1: Use derivative_label from file_info_map (from split)
            if file_path.name in filename_to_metadata:
                metadata = filename_to_metadata[file_path.name]
                derivative_label = metadata.get('derivative_label')
                if derivative_label and str(derivative_label).strip():
                    label = str(derivative_label).strip()
                    # Skip other label sources if we have derivative_label
            
            # Priority 2: Try common keys for labels in .npz file (only if derivative_label not found)
            if label is None:
                if 'label' in data:
                    label = data['label']
                elif 'labels' in data:
                    label = data['labels']
                elif 'y' in data:
                    label = data['y']
                elif 'keyword' in data:
                    label = data['keyword']
                elif 'keyword_id' in data:
                    label = data['keyword_id']
                elif 'keyword_name' in data:
                    label = data['keyword_name']
            
            # Priority 3: If no label in file, try to get from API using sample_id (fallback)
            if label is None and api_client and file_path.name in filename_to_metadata:
                metadata = filename_to_metadata[file_path.name]
                sample_id = metadata.get('sample_id')
                if sample_id:
                    try:
                        if sample_id not in keyword_cache:
                            # Get audio sample info
                            sample_info = api_client.audio.get_keyword_sample(sample_id)
                            keyword_id = sample_info.get('keyword_id')
                            if keyword_id:
                                # Get keyword info
                                keyword_info = api_client.keywords.get(keyword_id)
                                keyword_name = keyword_info.get('name', f'keyword_{keyword_id}')
                                keyword_cache[sample_id] = keyword_name
                            else:
                                keyword_cache[sample_id] = f'unknown_{sample_id}'
                        label = keyword_cache[sample_id]
                    except Exception as e:
                        print(f"Warning: Could not get keyword for sample {sample_id}: {e}")
                        label = f'unknown_{sample_id}'
            
            # Priority 4: Try to infer from filename (last resort)
            if label is None:
                label = file_path.stem.split('_')[0] if '_' in file_path.stem else file_path.stem
                print(f"Warning: Using filename as label: {label}")
            
            # Convert label to string if needed
            if isinstance(label, np.ndarray):
                label = label.item() if label.size == 1 else str(label)
            elif not isinstance(label, str):
                label = str(label)
            
            label_str = label.strip()
            
            # Map label to integer ID
            if label_str not in label_to_id:
                label_to_id[label_str] = current_id
                current_id += 1
                print(f"Added new label '{label_str}' with ID {label_to_id[label_str]}")
            
            label_id = label_to_id[label_str]
            
            # Handle different feature shapes
            if len(features.shape) == 1:
                # 1D features
                all_features.append(features)
            elif len(features.shape) == 2:
                # 2D features (e.g., time x frequency)
                all_features.append(features)
            elif len(features.shape) == 3:
                # 3D features - take first channel or flatten
                all_features.append(features[0] if features.shape[0] == 1 else features.flatten())
            else:
                # Flatten if needed
                all_features.append(features.flatten())
            
            all_labels.append(label_id)
            
        except Exception as e:
            print(f"Error loading {file_path}: {e}")
            import traceback
            traceback.print_exc()
            continue
    
    if not all_features:
        raise ValueError("No features loaded!")
    
    # Convert to numpy arrays
    # Find max shape for padding
    shapes = [f.shape for f in all_features]
    
    # If all have same shape, stack directly
    if len(set(shapes)) == 1:
        X = np.array(all_features)
    else:
        # Different shapes - need to handle padding or flattening
        # For CNN, we'll try to reshape to common format
        # Assume features are 2D (time, frequency) or 1D
        max_len = max([f.shape[0] if len(f.shape) >= 1 else len(f) for f in all_features])
        
        # Try to determine feature dimension
        if all(len(f.shape) == 2 for f in all_features):
            # All 2D - pad to same size
            max_freq = max([f.shape[1] for f in all_features])
            X = np.zeros((len(all_features), max_len, max_freq))
            for i, f in enumerate(all_features):
                X[i, :f.shape[0], :f.shape[1]] = f
        elif all(len(f.shape) == 1 for f in all_features):
            # All 1D - pad to same length
            X = np.zeros((len(all_features), max_len))
            for i, f in enumerate(all_features):
                X[i, :len(f)] = f
        else:
            # Mixed - flatten all
            max_size = max([f.size for f in all_features])
            X = np.zeros((len(all_features), max_size))
            for i, f in enumerate(all_features):
                X[i, :f.size] = f.flatten()
    
    y = np.array(all_labels)
    
    print(f"Loaded {len(X)} samples")
    print(f"Feature shape: {X.shape}")
    print(f"Number of classes: {len(label_to_id)}")
    print(f"Labels: {label_to_id}")
    
    return X, y, label_to_id