code/contrastive/generate_negatives_v4_FINAL.py

#!/usr/bin/env python3
"""
V5.1-FINAL: Corrected Negative Generation

FIXES from v3:
1. Sialic acid detection now uses WURCS pattern 'a2122h' (Neu5Ac) and 'a2212h' (Neu5Gc)
2. All token pattern matching updated for WURCS format
3. Rule5 (Sia extension) should now work correctly

OUTPUT: 100K rule-based negatives with proper biological patterns
"""

import sys
import pickle
import json
import random
import re
import numpy as np
from copy import deepcopy
from typing import Dict, List, Optional, Tuple
from tqdm import tqdm
import logging

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

# Paths
BASE = "/work/ratul1/supantha/glycan-SD-VS/bert_training_v3/v3.1_cluster_training"
BPE_VOCAB_PATH = f"{BASE}/data/bpe_vocabulary.json"
POSITIVES_PATH = f"{BASE}/bert_v5.1_contrastive/data/fully_resolved_161k.pkl"
OUTPUT_PATH = f"{BASE}/bert_v5.1_contrastive/data/hard_negatives_100k_v4_FINAL.pkl"

N_SAMPLES = 100000


class WURCSPatterns:
    """WURCS-specific patterns for monosaccharides and linkages."""
    
    # Monosaccharide patterns in WURCS
    SIALIC_ACID = ['a2122h', 'a2212h']  # Neu5Ac, Neu5Gc
    IDURONIC_ACID = ['a1221m']  # IdoA
    GLUCURONIC_ACID = ['a2112h']  # GlcA
    MANNOSE = ['a2112m']  # Man
    GLUCOSE = ['a2122m']  # Glc
    GALACTOSE = ['a2112a']  # Gal
    
    # Linkage patterns
    ALPHA_LINKS = ['-1a', 'a-']  # Alpha anomeric
    BETA_LINKS = ['-1b', 'b-']  # Beta anomeric


class UniversalBiologicalRulesV4:
    """
    V4: Proper WURCS pattern detection for all 8 rules.
    """
    
    def __init__(self, vocab_path: str, sequences: List[Dict]):
        logger.info(f"Loading BPE vocabulary from {vocab_path}")
        with open(vocab_path, 'r') as f:
            vocab = json.load(f)
        
        self.token_to_id = vocab['token_to_id']
        self.id_to_token = {v: k for k, v in self.token_to_id.items()}
        self.vocab_size = len(self.token_to_id)
        self.sequences = sequences
        
        # Categorize tokens by WURCS patterns
        self._categorize_tokens()
        
    def _categorize_tokens(self):
        """Categorize vocabulary tokens by WURCS patterns."""
        self.sialic_tokens = []
        self.idoa_tokens = []
        self.alpha_linkage_tokens = []
        self.beta_linkage_tokens = []
        self.monosaccharide_tokens = []
        
        for token, tid in self.token_to_id.items():
            # Sialic acid (Neu5Ac = a2122h, Neu5Gc = a2212h)
            if any(pattern in token for pattern in WURCSPatterns.SIALIC_ACID):
                self.sialic_tokens.append(tid)
            
            # Iduronic acid (IdoA = a1221m)
            if any(pattern in token for pattern in WURCSPatterns.IDURONIC_ACID):
                self.idoa_tokens.append(tid)
            
            # Alpha linkages (-1a or ends with 'a')
            if any(pattern in token for pattern in WURCSPatterns.ALPHA_LINKS):
                self.alpha_linkage_tokens.append(tid)
            
            # Beta linkages (-1b or ends with 'b')
            if any(pattern in token for pattern in WURCSPatterns.BETA_LINKS):
                self.beta_linkage_tokens.append(tid)
            
            # General monosaccharides (start with 'a' followed by digits)
            if re.match(r'^a\d+[a-z]', token):
                self.monosaccharide_tokens.append(tid)
        
        logger.info(f"WURCS Token Categories:")
        logger.info(f"  Sialic acid (a2122h/a2212h): {len(self.sialic_tokens)} tokens")
        logger.info(f"  IdoA (a1221m): {len(self.idoa_tokens)} tokens")
        logger.info(f"  Alpha linkages: {len(self.alpha_linkage_tokens)} tokens")
        logger.info(f"  Beta linkages: {len(self.beta_linkage_tokens)} tokens")
        logger.info(f"  Monosaccharides: {len(self.monosaccharide_tokens)} tokens")
    
    def rule1_sia_beta_linkage(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 1: Sialic acid with β-linkage (sialyltransferases only produce α)."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        # Find sialic acid positions
        sia_positions = [i for i, tid in enumerate(token_ids) if tid in self.sialic_tokens]
        
        if sia_positions and self.beta_linkage_tokens:
            # Change the linkage near Sia to beta
            pos = random.choice(sia_positions)
            if pos + 1 < len(token_ids):
                token_ids[pos + 1] = random.choice(self.beta_linkage_tokens)
        elif self.sialic_tokens and self.beta_linkage_tokens and len(token_ids) > 10:
            # Insert Sia with beta linkage
            pos = random.randint(3, len(token_ids) - 3)
            token_ids.insert(pos, random.choice(self.sialic_tokens))
            token_ids.insert(pos + 1, random.choice(self.beta_linkage_tokens))
        else:
            return None
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule1_sia_beta_linkage'
        return mutated
    
    def rule2_idoa_donor(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 2: IdoA as glycosyl donor (only formed by epimerization)."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 8:
            return None
        
        # Add IdoA at a random position as if it were a donor
        pos = random.randint(2, len(token_ids) - 2)
        if self.idoa_tokens:
            token_ids.insert(pos, random.choice(self.idoa_tokens))
        else:
            # Create invalid pattern
            token_ids[pos] = random.randint(100, self.vocab_size - 1)
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule2_idoa_donor'
        return mutated
    
    def rule3_high_linkage(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 3: Linkage positions ≥8 (except α2-8 polySia)."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 8:
            return None
        
        # Insert impossible high linkage number
        pos = random.randint(2, len(token_ids) - 2)
        # Use digit tokens for 8, 9
        high_digits = [tid for t, tid in self.token_to_id.items() if t in ['8', '9']]
        if high_digits:
            token_ids.insert(pos, random.choice(high_digits))
            token_ids.insert(pos + 1, random.choice(high_digits))
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule3_high_linkage'
        return mutated
    
    def rule4_overbranching(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 4: More than 4 branches per residue."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 10:
            return None
        
        # Add multiple branch tokens
        branch_open = self.token_to_id.get('[BRANCH_OPEN]', 5)
        branch_close = self.token_to_id.get('[BRANCH_CLOSE]', 6)
        
        pos = random.randint(3, len(token_ids) - 5)
        for _ in range(5):  # 5 branches = impossible
            token_ids.insert(pos, branch_open)
        for _ in range(5):
            token_ids.append(branch_close)
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule4_overbranching'
        return mutated
    
    def rule5_sia_extension(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 5: Extending past sialic acid (no GT accepts Sia as acceptor)."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        # Find sialic acid
        sia_positions = [i for i, tid in enumerate(token_ids) if tid in self.sialic_tokens]
        
        if sia_positions:
            # Add sugar AFTER sialic acid (impossible)
            pos = random.choice(sia_positions)
            if self.monosaccharide_tokens and pos + 2 < len(token_ids):
                # Insert a monosaccharide after Sia
                token_ids.insert(pos + 1, random.choice(self.monosaccharide_tokens))
                if self.alpha_linkage_tokens:
                    token_ids.insert(pos + 2, random.choice(self.alpha_linkage_tokens))
        else:
            # No Sia found - add Sia then extend it
            if self.sialic_tokens and self.monosaccharide_tokens and len(token_ids) > 10:
                pos = random.randint(5, len(token_ids) - 5)
                token_ids.insert(pos, random.choice(self.sialic_tokens))
                token_ids.insert(pos + 1, random.choice(self.monosaccharide_tokens))
            else:
                return None
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule5_sia_extension'
        return mutated
    
    def rule6_c5_linkage(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 6: C5 ring oxygen linkage (impossible)."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 8:
            return None
        
        # Insert "5" as linkage position
        five_token = self.token_to_id.get('5', 24)
        pos = random.randint(2, len(token_ids) - 2)
        token_ids.insert(pos, five_token)
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule6_c5_linkage'
        return mutated
    
    def rule7_multi_anomeric(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 7: Multiple anomeric bonds from same carbon."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 8:
            return None
        
        # Add conflicting anomeric markers
        pos = random.randint(2, len(token_ids) - 2)
        if self.alpha_linkage_tokens and self.beta_linkage_tokens:
            token_ids.insert(pos, random.choice(self.alpha_linkage_tokens))
            token_ids.insert(pos + 1, random.choice(self.beta_linkage_tokens))
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule7_multi_anomeric'
        return mutated
    
    def rule8_c1_acetal(self, seq_data: Dict) -> Optional[Dict]:
        """Rule 8: Free reducing end conflict with glycosidic bond."""
        mutated = deepcopy(seq_data)
        token_ids = list(mutated.get('token_ids', []))
        
        if len(token_ids) < 8:
            return None
        
        # Add anomeric configuration at end
        pos = len(token_ids) - 2
        if self.alpha_linkage_tokens:
            token_ids.insert(pos, random.choice(self.alpha_linkage_tokens))
        
        mutated['token_ids'] = token_ids[:256]
        mutated['is_negative'] = True
        mutated['negative_method'] = 'rule8_c1_acetal'
        return mutated
    
    def generate_negatives(self, n_samples: int = 100000) -> Tuple[List[Dict], Dict]:
        """Generate negatives with equal distribution across all 8 rules."""
        logger.info(f"Generating {n_samples} negatives with 8 Universal Rules...")
        
        rules = [
            (self.rule1_sia_beta_linkage, 'rule1_sia_beta'),
            (self.rule2_idoa_donor, 'rule2_idoa_donor'),
            (self.rule3_high_linkage, 'rule3_high_linkage'),
            (self.rule4_overbranching, 'rule4_overbranching'),
            (self.rule5_sia_extension, 'rule5_sia_extension'),
            (self.rule6_c5_linkage, 'rule6_c5_linkage'),
            (self.rule7_multi_anomeric, 'rule7_multi_anomeric'),
            (self.rule8_c1_acetal, 'rule8_c1_acetal'),
        ]
        
        per_rule = n_samples // len(rules)
        negatives = []
        stats = {name: 0 for _, name in rules}
        
        pbar = tqdm(total=n_samples, desc="Generating 8-rule negatives")
        
        for rule_fn, rule_name in rules:
            count = 0
            attempts = 0
            max_attempts = per_rule * 5
            
            while count < per_rule and attempts < max_attempts:
                source = random.choice(self.sequences)
                neg = rule_fn(source)
                if neg:
                    negatives.append(neg)
                    stats[rule_name] += 1
                    count += 1
                    pbar.update(1)
                attempts += 1
        
        pbar.close()
        logger.info(f"Generated {len(negatives)} negatives")
        logger.info(f"Distribution: {json.dumps(stats, indent=2)}")
        
        return negatives, stats


def main():
    print("=" * 70)
    print("V5.1-FINAL: Corrected 8 Rules with WURCS Patterns")
    print("=" * 70)
    
    logger.info("Loading positive sequences...")
    with open(POSITIVES_PATH, 'rb') as f:
        positives = pickle.load(f)
    logger.info(f"Loaded {len(positives)} positives")
    
    generator = UniversalBiologicalRulesV4(BPE_VOCAB_PATH, positives)
    negatives, stats = generator.generate_negatives(n_samples=N_SAMPLES)
    
    logger.info(f"Saving {len(negatives)} negatives...")
    with open(OUTPUT_PATH, 'wb') as f:
        pickle.dump(negatives, f)
    
    # Save stats
    stats_output = {
        "total_negatives": len(negatives),
        "source_positives": len(positives),
        "mutation_types": stats,
        "version": "v4_FINAL_WURCS",
        "fixes": [
            "Sialic acid detection now uses WURCS pattern a2122h/a2212h",
            "All token matching updated for WURCS format"
        ]
    }
    
    stats_path = OUTPUT_PATH.replace('.pkl', '_stats.json')
    with open(stats_path, 'w') as f:
        json.dump(stats_output, f, indent=2)
    
    print()
    print("=" * 70)
    print(f"Generated {len(negatives):,} rule-based negatives")
    print(f"Saved to: {OUTPUT_PATH}")
    print("=" * 70)
    for rule, count in sorted(stats.items()):
        pct = 100 * count / len(negatives) if negatives else 0
        print(f"  {rule}: {count:,} ({pct:.1f}%)")


if __name__ == '__main__':
    main()