src/cfg_dataset.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import numpy as np
import json
import os
from typing import Dict, List, Tuple, Optional
import random
import re

def parse_fasta_with_amp_labels(fasta_path: str, max_seq_len: int = 50) -> Dict[str, any]:
    """
    Parse FASTA file and assign AMP/Non-AMP labels based on header prefixes.
    
    Label assignment strategy:
    - AMP (0): Headers starting with '>AP'
    - Non-AMP (1): Headers starting with '>sp'
    - Mask (2): Used for CFG training (randomly assigned)
    
    File format:
    - Odd lines: Headers (>sp or >AP)
    - Even lines: Amino acid sequences
    
    Args:
        fasta_path: Path to FASTA file
        max_seq_len: Maximum sequence length to include
        
    Returns:
        Dictionary with sequences, labels, and metadata
    """
    sequences = []
    labels = []
    headers = []
    
    print(f"Parsing FASTA file: {fasta_path}")
    print("Label assignment: >AP = AMP (0), >sp = Non-AMP (1)")
    
    current_header = ""
    current_sequence = ""
    
    with open(fasta_path, 'r') as f:
        for line in f:
            line = line.strip()
            if line.startswith('>'):
                # Process previous sequence if exists
                if current_sequence and current_header:
                    if 2 <= len(current_sequence) <= max_seq_len:
                        # Check if sequence contains only canonical amino acids
                        canonical_aa = set('ACDEFGHIKLMNPQRSTVWY')
                        if all(aa in canonical_aa for aa in current_sequence.upper()):
                            sequences.append(current_sequence.upper())
                            headers.append(current_header)
                            
                            # Assign label based on header prefix
                            if current_header.startswith('AP'):
                                labels.append(0)  # AMP
                            elif current_header.startswith('sp'):
                                labels.append(1)  # Non-AMP
                            else:
                                # Unknown prefix, default to Non-AMP
                                labels.append(1)
                                print(f"Warning: Unknown header prefix in '{current_header}', defaulting to Non-AMP")
                
                # Start new sequence (remove '>' from header)
                current_header = line[1:]
                current_sequence = ""
            else:
                current_sequence += line
        
        # Process last sequence
        if current_sequence and current_header:
            if 2 <= len(current_sequence) <= max_seq_len:
                canonical_aa = set('ACDEFGHIKLMNPQRSTVWY')
                if all(aa in canonical_aa for aa in current_sequence.upper()):
                    sequences.append(current_sequence.upper())
                    headers.append(current_header)
                    
                    # Assign label based on header prefix
                    if current_header.startswith('AP'):
                        labels.append(0)  # AMP
                    elif current_header.startswith('sp'):
                        labels.append(1)  # Non-AMP
                    else:
                        # Unknown prefix, default to Non-AMP
                        labels.append(1)
                        print(f"Warning: Unknown header prefix in '{current_header}', defaulting to Non-AMP")
    
    # Create masked labels for CFG training (10% masked)
    original_labels = np.array(labels)
    masked_labels = original_labels.copy()
    mask_probability = 0.1
    mask_indices = np.random.choice(
        len(original_labels), 
        size=int(len(original_labels) * mask_probability), 
        replace=False
    )
    masked_labels[mask_indices] = 2  # 2 = mask/unknown
    
    print(f"✓ Parsed {len(sequences)} valid sequences from FASTA")
    print(f"  AMP sequences: {np.sum(original_labels == 0)}")
    print(f"  Non-AMP sequences: {np.sum(original_labels == 1)}")
    print(f"  Masked for CFG: {len(mask_indices)}")
    
    return {
        'sequences': sequences,
        'headers': headers,
        'labels': original_labels.tolist(),
        'masked_labels': masked_labels.tolist(),
        'mask_indices': mask_indices.tolist()
    }

class CFGUniProtDataset(Dataset):
    """
    Dataset class for UniProt sequences with classifier-free guidance.
    
    This dataset:
    1. Loads processed UniProt data with AMP classifications
    2. Handles label masking for CFG training
    3. Integrates with your existing flow training pipeline
    4. Provides sequences, labels, and masking information
    """
    
    def __init__(self, 
                 data_path: str,
                 use_masked_labels: bool = True,
                 mask_probability: float = 0.1,
                 max_seq_len: int = 50,
                 device: str = 'cuda'):
        
        self.data_path = data_path
        self.use_masked_labels = use_masked_labels
        self.mask_probability = mask_probability
        self.max_seq_len = max_seq_len
        self.device = device
        
        # Load processed data
        self._load_data()
        
        # Label mapping
        self.label_map = {
            0: 'amp',      # MIC < 100
            1: 'non_amp',  # MIC > 100
            2: 'mask'      # Unknown MIC
        }
        
        print(f"CFG Dataset initialized:")
        print(f"  Total sequences: {len(self.sequences)}")
        print(f"  Using masked labels: {use_masked_labels}")
        print(f"  Mask probability: {mask_probability}")
        print(f"  Label distribution: {self._get_label_distribution()}")
    
    def _load_data(self):
        """Load processed UniProt data."""
        if os.path.exists(self.data_path):
            with open(self.data_path, 'r') as f:
                data = json.load(f)
            
            self.sequences = data['sequences']
            self.original_labels = np.array(data['original_labels'])
            self.masked_labels = np.array(data['masked_labels'])
            self.mask_indices = set(data['mask_indices'])
            
        else:
            raise FileNotFoundError(f"Data file not found: {self.data_path}")
    
    def _get_label_distribution(self) -> Dict[str, int]:
        """Get distribution of labels in the dataset."""
        labels = self.masked_labels if self.use_masked_labels else self.original_labels
        unique, counts = np.unique(labels, return_counts=True)
        return {self.label_map[label]: count for label, count in zip(unique, counts)}
    
    def __len__(self) -> int:
        return len(self.sequences)
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get a single sample with sequence and label."""
        sequence = self.sequences[idx]
        
        # Get appropriate label
        if self.use_masked_labels:
            label = self.masked_labels[idx]
        else:
            label = self.original_labels[idx]
        
        # Check if this sample was masked
        is_masked = idx in self.mask_indices
        
        return {
            'sequence': sequence,
            'label': torch.tensor(label, dtype=torch.long),
            'original_label': torch.tensor(self.original_labels[idx], dtype=torch.long),
            'is_masked': torch.tensor(is_masked, dtype=torch.bool),
            'index': torch.tensor(idx, dtype=torch.long)
        }
    
    def get_label_statistics(self) -> Dict[str, Dict]:
        """Get detailed statistics about labels."""
        stats = {
            'original': self._get_label_distribution(),
            'masked': self._get_label_distribution() if self.use_masked_labels else None,
            'masking_info': {
                'total_masked': len(self.mask_indices),
                'mask_probability': self.mask_probability,
                'masked_indices': list(self.mask_indices)
            }
        }
        return stats

class CFGFlowDataset(Dataset):
    """
    Dataset that integrates CFG labels with your existing flow training pipeline.
    
    This dataset:
    1. Loads your existing AMP embeddings
    2. Adds CFG labels from UniProt processing
    3. Handles the integration between embeddings and labels
    4. Provides data in the format expected by your flow training
    """
    
    def __init__(self, 
                 embeddings_path: str,
                 cfg_data_path: str,
                 use_masked_labels: bool = True,
                 max_seq_len: int = 50,
                 device: str = 'cuda'):
        
        self.embeddings_path = embeddings_path
        self.cfg_data_path = cfg_data_path
        self.use_masked_labels = use_masked_labels
        self.max_seq_len = max_seq_len
        self.device = device
        
        # Load data
        self._load_embeddings()
        self._load_cfg_data()
        self._align_data()
        
        print(f"CFG Flow Dataset initialized:")
        print(f"  AMP embeddings: {self.embeddings.shape}")
        print(f"  CFG labels: {len(self.cfg_labels)}")
        print(f"  Aligned samples: {len(self.aligned_indices)}")
    
    def _load_embeddings(self):
        """Load your existing AMP embeddings."""
        print(f"Loading AMP embeddings from {self.embeddings_path}...")
        
        # Try to load the combined embeddings file first (FULL DATA)
        combined_path = os.path.join(self.embeddings_path, "all_peptide_embeddings.pt")
        
        if os.path.exists(combined_path):
            print(f"Loading combined embeddings from {combined_path} (FULL DATA)...")
            # Load on CPU first to avoid CUDA issues with DataLoader workers
            self.embeddings = torch.load(combined_path, map_location='cpu')
            print(f"✓ Loaded ALL embeddings: {self.embeddings.shape}")
        else:
            print("Combined embeddings file not found, loading individual files...")
            # Fallback to individual files
            import glob
            
            embedding_files = glob.glob(os.path.join(self.embeddings_path, "*.pt"))
            embedding_files = [f for f in embedding_files if not f.endswith('metadata.json') and not f.endswith('sequence_ids.json') and not f.endswith('all_peptide_embeddings.pt')]
            
            print(f"Found {len(embedding_files)} individual embedding files")
            
            # Load and stack all embeddings
            embeddings_list = []
            for file_path in embedding_files:
                try:
                    embedding = torch.load(file_path, map_location='cpu')
                    if embedding.dim() == 2:  # (seq_len, hidden_dim)
                        embeddings_list.append(embedding)
                    else:
                        print(f"Warning: Skipping {file_path} - unexpected shape {embedding.shape}")
                except Exception as e:
                    print(f"Warning: Could not load {file_path}: {e}")
            
            if not embeddings_list:
                raise ValueError("No valid embeddings found!")
            
            self.embeddings = torch.stack(embeddings_list)
            print(f"Loaded {len(self.embeddings)} embeddings from individual files")
    
    def _load_cfg_data(self):
        """Load CFG data from FASTA file with automatic AMP labeling."""
        print(f"Loading CFG data from FASTA: {self.cfg_data_path}...")
        
        # Check if it's a FASTA file or JSON file
        if self.cfg_data_path.endswith('.fasta') or self.cfg_data_path.endswith('.fa'):
            # Parse FASTA file with automatic labeling
            cfg_data = parse_fasta_with_amp_labels(self.cfg_data_path, self.max_seq_len)
            
            self.cfg_sequences = cfg_data['sequences']
            self.cfg_headers = cfg_data['headers']
            self.cfg_original_labels = np.array(cfg_data['labels'])
            self.cfg_masked_labels = np.array(cfg_data['masked_labels'])
            self.cfg_mask_indices = set(cfg_data['mask_indices'])
            
        else:
            # Legacy JSON format support
            with open(self.cfg_data_path, 'r') as f:
                cfg_data = json.load(f)
            
            self.cfg_sequences = cfg_data['sequences']
            self.cfg_headers = cfg_data.get('headers', [''] * len(cfg_data['sequences']))
            self.cfg_original_labels = np.array(cfg_data['labels'])
            
            # For CFG training, we need to create masked labels
            # Randomly mask 10% of labels for CFG training
            self.cfg_masked_labels = self.cfg_original_labels.copy()
            mask_probability = 0.1
            mask_indices = np.random.choice(
                len(self.cfg_original_labels), 
                size=int(len(self.cfg_original_labels) * mask_probability), 
                replace=False
            )
            self.cfg_masked_labels[mask_indices] = 2  # 2 = mask/unknown
            self.cfg_mask_indices = set(mask_indices)
        
        print(f"Loaded {len(self.cfg_sequences)} CFG sequences")
        print(f"Label distribution: {np.bincount(self.cfg_original_labels)}")
        print(f"Masked {len(self.cfg_mask_indices)} labels for CFG training")
    
    def _align_data(self):
        """Align AMP embeddings with CFG data based on sequence matching."""
        print("Aligning AMP embeddings with CFG data...")
        
        # For now, we'll use a simple approach: take the first N sequences
        # where N is the minimum of embeddings and CFG data
        min_samples = min(len(self.embeddings), len(self.cfg_sequences))
        
        self.aligned_indices = list(range(min_samples))
        
        # Align labels
        if self.use_masked_labels:
            self.cfg_labels = self.cfg_masked_labels[:min_samples]
        else:
            self.cfg_labels = self.cfg_original_labels[:min_samples]
        
        # Align embeddings
        self.aligned_embeddings = self.embeddings[:min_samples]
        
        print(f"Aligned {min_samples} samples")
    
    def __len__(self) -> int:
        return len(self.aligned_indices)
    
    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        """Get a single sample with embedding and CFG label."""
        # Embeddings are already on CPU
        embedding = self.aligned_embeddings[idx]
        label = self.cfg_labels[idx]
        original_label = self.cfg_original_labels[idx]
        is_masked = idx in self.cfg_mask_indices
        
        return {
            'embedding': embedding,
            'label': torch.tensor(label, dtype=torch.long),
            'original_label': torch.tensor(original_label, dtype=torch.long),
            'is_masked': torch.tensor(is_masked, dtype=torch.bool),
            'index': torch.tensor(idx, dtype=torch.long)
        }
    
    def get_embedding_stats(self) -> Dict:
        """Get statistics about the embeddings."""
        return {
            'shape': self.aligned_embeddings.shape,
            'mean': self.aligned_embeddings.mean().item(),
            'std': self.aligned_embeddings.std().item(),
            'min': self.aligned_embeddings.min().item(),
            'max': self.aligned_embeddings.max().item()
        }

def create_cfg_dataloader(dataset: Dataset, 
                         batch_size: int = 32,
                         shuffle: bool = True,
                         num_workers: int = 4) -> DataLoader:
    """Create a DataLoader for CFG training."""
    
    def collate_fn(batch):
        """Custom collate function for CFG data."""
        # Separate different types of data
        embeddings = torch.stack([item['embedding'] for item in batch])
        labels = torch.stack([item['label'] for item in batch])
        original_labels = torch.stack([item['original_label'] for item in batch])
        is_masked = torch.stack([item['is_masked'] for item in batch])
        indices = torch.stack([item['index'] for item in batch])
        
        return {
            'embeddings': embeddings,
            'labels': labels,
            'original_labels': original_labels,
            'is_masked': is_masked,
            'indices': indices
        }
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=True
    )

def test_cfg_dataset():
    """Test function to verify the CFG dataset works correctly."""
    print("Testing CFG Dataset...")
    
    # Test with a small subset
    test_data = {
        'sequences': ['MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG', 
                     'MKLLIVTFCLTFAAL',
                     'MKLLIVTFCLTFAALMKLLIVTFCLTFAAL'],
        'original_labels': [0, 1, 0],  # amp, non_amp, amp
        'masked_labels': [0, 2, 0],    # amp, mask, amp
        'mask_indices': [1]  # Only second sequence is masked
    }
    
    # Save test data
    test_path = 'test_cfg_data.json'
    with open(test_path, 'w') as f:
        json.dump(test_data, f)
    
    # Test dataset
    dataset = CFGUniProtDataset(test_path, use_masked_labels=True)
    
    print(f"Dataset length: {len(dataset)}")
    for i in range(len(dataset)):
        sample = dataset[i]
        print(f"Sample {i}:")
        print(f"  Sequence: {sample['sequence'][:20]}...")
        print(f"  Label: {sample['label'].item()}")
        print(f"  Original Label: {sample['original_label'].item()}")
        print(f"  Is Masked: {sample['is_masked'].item()}")
    
    # Clean up
    os.remove(test_path)
    print("Test completed successfully!")

if __name__ == "__main__":
    test_cfg_dataset()