model.py

"""
Optimized GraphCodeBERT Model for JavaScript Vulnerability Detection
With comprehensive training tracking and model saving
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from transformers import RobertaModel, AdamW, get_linear_schedule_with_warmup
import numpy as np
import pandas as pd
from tqdm import tqdm
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
import os
import json
from datetime import datetime

# CPU Optimization
torch.set_num_threads(os.cpu_count())
if hasattr(torch, 'set_float32_matmul_precision'):
    torch.set_float32_matmul_precision('medium')

from __init__ import GraphCodeBERTProcessor


class VulnerabilityDataset(Dataset):
    """Dataset for vulnerability detection with code and numeric features"""
    
    def __init__(self, codes, labels, numeric_features=None, max_length=512, model_name='microsoft/graphcodebert-base'):
        self.codes = codes
        self.labels = labels
        self.numeric_features = numeric_features
        self.max_length = max_length
        self.model_name = model_name
        self.processor = None
    
    def _get_processor(self):
        if self.processor is None:
            self.processor = GraphCodeBERTProcessor(
                model_name=self.model_name,
                max_length=self.max_length
            )
        return self.processor
    
    def __len__(self):
        return len(self.codes)
    
    def __getitem__(self, idx):
        code = self.codes[idx]
        label = self.labels[idx]
        
        processor = self._get_processor()
        processed = processor.process_code(code)
        
        item = {
            'input_ids': processed['input_ids'].squeeze(0),
            'attention_mask': processed['attention_mask'].squeeze(0),
            'adjacency_matrix': processed['adjacency_matrix'],
            'label': torch.tensor(label, dtype=torch.long)
        }
        
        if self.numeric_features is not None:
            item['numeric_features'] = torch.tensor(
                self.numeric_features[idx], 
                dtype=torch.float32
            )
        
        return item


class FocalLoss(nn.Module):
    """Focal Loss for handling class imbalance"""
    def __init__(self, alpha=0.25, gamma=2.0):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
    
    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class GraphAttentionLayer(nn.Module):
    """Optimized Graph Attention Layer"""
    
    def __init__(self, in_features, out_features, dropout=0.1, alpha=0.2, simplified=True):
        super(GraphAttentionLayer, self).__init__()
        self.simplified = simplified
        self.dropout = dropout
        
        if simplified:
            self.W = nn.Linear(in_features, out_features, bias=False)
        else:
            self.W = nn.Parameter(torch.zeros(size=(in_features, out_features)))
            nn.init.xavier_uniform_(self.W.data, gain=1.414)
            self.a = nn.Parameter(torch.zeros(size=(2 * out_features, 1)))
            nn.init.xavier_uniform_(self.a.data, gain=1.414)
            self.leakyrelu = nn.LeakyReLU(alpha)
    
    def forward(self, h, adj):
        if self.simplified:
            Wh = self.W(h)
            degree = adj.sum(dim=-1, keepdim=True).clamp(min=1)
            adj_norm = adj / degree
            h_prime = torch.matmul(adj_norm, Wh)
            return F.dropout(h_prime, self.dropout, training=self.training)
        else:
            Wh = torch.matmul(h, self.W)
            batch_size, seq_len, out_features = Wh.shape
            
            attention_scores = []
            for i in range(batch_size):
                Wh_i = Wh[i]
                Wh1 = Wh_i.unsqueeze(1).expand(seq_len, seq_len, out_features)
                Wh2 = Wh_i.unsqueeze(0).expand(seq_len, seq_len, out_features)
                combined = torch.cat([Wh1, Wh2], dim=-1)
                e = self.leakyrelu(torch.matmul(combined, self.a).squeeze(-1))
                attention_scores.append(e)
            
            e = torch.stack(attention_scores, dim=0)
            zero_vec = -9e15 * torch.ones_like(e)
            attention = torch.where(adj > 0, e, zero_vec)
            attention = F.softmax(attention, dim=-1)
            attention = F.dropout(attention, self.dropout, training=self.training)
            h_prime = torch.matmul(attention, Wh)
            return h_prime


class GraphCodeBERTModel(nn.Module):
    """Enhanced GraphCodeBERT with Graph Attention - PROPERLY USING ADJACENCY MATRIX"""
    
    def __init__(
        self, 
        model_name='microsoft/graphcodebert-base',
        num_labels=2,
        numeric_feature_dim=0,
        hidden_dim=768,
        dropout=0.2,
        num_gat_layers=2,
        use_simplified_gat=True
    ):
        super(GraphCodeBERTModel, self).__init__()
        
        self.graphcodebert = RobertaModel.from_pretrained(model_name)
        self.config = self.graphcodebert.config
        
        # GAT Layers with residual connections - THESE USE THE ADJACENCY MATRIX
        self.gat_layers = nn.ModuleList([
            GraphAttentionLayer(hidden_dim, hidden_dim, dropout, simplified=use_simplified_gat)
            for _ in range(num_gat_layers)
        ])
        
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(hidden_dim) for _ in range(num_gat_layers)
        ])
        
        # Feature fusion
        self.numeric_feature_dim = numeric_feature_dim
        if numeric_feature_dim > 0:
            self.numeric_projector = nn.Sequential(
                nn.Linear(numeric_feature_dim, 128),
                nn.BatchNorm1d(128),
                nn.ReLU(),
                nn.Dropout(dropout),
                nn.Linear(128, 64),
                nn.ReLU()
            )
            fusion_dim = hidden_dim + 64
        else:
            fusion_dim = hidden_dim
        
        # Enhanced classification head
        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(dropout * 0.5),
            nn.Linear(128, num_labels)
        )
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, input_ids, attention_mask, adjacency_matrix, numeric_features=None):
        """
        Forward pass - ADJACENCY MATRIX IS ACTIVELY USED HERE
        
        The adjacency_matrix encodes:
        - Value 1: AST edges (parent-child relationships in syntax tree)
        - Value 2: DFG edges (data flow between variables)
        - Value 0: No connection
        
        GAT layers use this matrix to aggregate information from connected nodes
        """
        # Get GraphCodeBERT embeddings
        outputs = self.graphcodebert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        
        sequence_output = outputs.last_hidden_state  # (batch, seq_len, hidden_dim)
        
        # Convert adjacency matrix to binary (any edge = 1)
        # This allows GAT to process both AST and DFG edges uniformly
        adj_binary = (adjacency_matrix > 0).float()
        
        # Apply GAT layers with residual connections
        # CRITICAL: adjacency_matrix guides how nodes aggregate information
        graph_output = sequence_output
        for gat_layer, layer_norm in zip(self.gat_layers, self.layer_norms):
            graph_out = gat_layer(graph_output, adj_binary)  # <-- ADJACENCY MATRIX USED HERE
            graph_output = layer_norm(graph_output + graph_out)  # Residual connection
            graph_output = self.dropout(graph_output)
        
        # Pool: take [CLS] token representation
        pooled_output = graph_output[:, 0, :]
        
        # Fuse with numeric features if available
        if numeric_features is not None and self.numeric_feature_dim > 0:
            numeric_embed = self.numeric_projector(numeric_features)
            pooled_output = torch.cat([pooled_output, numeric_embed], dim=-1)
        
        # Classification
        logits = self.classifier(pooled_output)
        return logits


class VulnerabilityDetector:
    """Main training and evaluation class with comprehensive tracking"""
    
    def __init__(
        self,
        model_name='microsoft/graphcodebert-base',
        numeric_feature_dim=0,
        device=None,
        learning_rate=2e-5,
        weight_decay=0.01,
        num_epochs=15,  # Increased for better convergence
        batch_size=16,
        max_length=512,
        patience=5  # Early stopping patience
    ):
        if device is None:
            if torch.cuda.is_available():
                self.device = torch.device('cuda')
                print(f"✓ Using GPU: {torch.cuda.get_device_name(0)}")
                print(f"  GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
            else:
                self.device = torch.device('cpu')
                print("✓ Using CPU")
        else:
            self.device = device
        
        self.learning_rate = learning_rate
        self.weight_decay = weight_decay
        self.num_epochs = num_epochs
        self.batch_size = batch_size
        self.max_length = max_length
        self.patience = patience
        
        self.model = GraphCodeBERTModel(
            model_name=model_name,
            num_labels=2,
            numeric_feature_dim=numeric_feature_dim
        ).to(self.device)
        
        print(f"✓ Model initialized on {self.device}")
        print(f"✓ Total parameters: {sum(p.numel() for p in self.model.parameters()):,}")
        print(f"✓ Trainable parameters: {sum(p.numel() for p in self.model.parameters() if p.requires_grad):,}")
    
    def prepare_data_loaders(self, train_data, val_data, test_data):
        """Prepare PyTorch data loaders"""
        X_train, y_train, train_numeric = train_data
        X_val, y_val, val_numeric = val_data
        X_test, y_test, test_numeric = test_data
        
        train_dataset = VulnerabilityDataset(
            X_train['code'].values, y_train.values, train_numeric,
            max_length=self.max_length, model_name='microsoft/graphcodebert-base'
        )
        
        val_dataset = VulnerabilityDataset(
            X_val['code'].values, y_val.values, val_numeric,
            max_length=self.max_length, model_name='microsoft/graphcodebert-base'
        )
        
        test_dataset = VulnerabilityDataset(
            X_test['code'].values, y_test.values, test_numeric,
            max_length=self.max_length, model_name='microsoft/graphcodebert-base'
        )
        
        train_loader = DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=0)
        val_loader = DataLoader(val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=0)
        test_loader = DataLoader(test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=0)
        
        return train_loader, val_loader, test_loader
    
    def train(self, train_loader, val_loader):
        """Train with comprehensive tracking"""
        
        # Optimizer with weight decay
        optimizer = AdamW(
            self.model.parameters(),
            lr=self.learning_rate,
            weight_decay=self.weight_decay
        )
        
        total_steps = len(train_loader) * self.num_epochs
        scheduler = get_linear_schedule_with_warmup(
            optimizer,
            num_warmup_steps=int(total_steps * 0.1),
            num_training_steps=total_steps
        )
        
        # Use Focal Loss for class imbalance
        criterion = FocalLoss(alpha=0.25, gamma=2.0)
        
        # Mixed precision
        use_amp = self.device.type == 'cuda'
        scaler = torch.cuda.amp.GradScaler() if use_amp else None
        
        # Tracking
        history = {
            'train_loss': [], 'train_acc': [], 'train_f1': [],
            'val_loss': [], 'val_acc': [], 'val_f1': [],
            'learning_rates': []
        }
        
        best_val_f1 = 0
        patience_counter = 0
        
        print("\n" + "="*60)
        print("STARTING TRAINING")
        if use_amp:
            print("✓ Using mixed precision (FP16)")
        print("="*60)
        
        for epoch in range(self.num_epochs):
            # Training
            self.model.train()
            train_loss = 0
            train_preds, train_labels = [], []
            
            train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs} [Train]")
            
            for batch in train_pbar:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                adjacency_matrix = batch['adjacency_matrix'].to(self.device)
                labels = batch['label'].to(self.device)
                numeric_features = batch.get('numeric_features', None)
                if numeric_features is not None:
                    numeric_features = numeric_features.to(self.device)
                
                optimizer.zero_grad()
                
                if use_amp:
                    with torch.cuda.amp.autocast():
                        logits = self.model(input_ids, attention_mask, adjacency_matrix, numeric_features)
                        loss = criterion(logits, labels)
                    scaler.scale(loss).backward()
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    logits = self.model(input_ids, attention_mask, adjacency_matrix, numeric_features)
                    loss = criterion(logits, labels)
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    optimizer.step()
                
                scheduler.step()
                
                train_loss += loss.item()
                preds = torch.argmax(logits, dim=-1)
                train_preds.extend(preds.cpu().numpy())
                train_labels.extend(labels.cpu().numpy())
                
                train_pbar.set_postfix({'loss': f'{loss.item():.4f}'})
            
            avg_train_loss = train_loss / len(train_loader)
            train_acc = accuracy_score(train_labels, train_preds)
            train_f1 = f1_score(train_labels, train_preds, zero_division=0)
            
            # Validation
            val_metrics = self.evaluate(val_loader, criterion)
            
            # Track history
            history['train_loss'].append(avg_train_loss)
            history['train_acc'].append(train_acc)
            history['train_f1'].append(train_f1)
            history['val_loss'].append(val_metrics['loss'])
            history['val_acc'].append(val_metrics['accuracy'])
            history['val_f1'].append(val_metrics['f1'])
            history['learning_rates'].append(scheduler.get_last_lr()[0])
            
            # Print epoch results
            print(f"\n{'='*60}")
            print(f"Epoch {epoch+1}/{self.num_epochs}")
            print(f"{'='*60}")
            print(f"Train Loss: {avg_train_loss:.4f} | Train Acc: {train_acc:.4f} | Train F1: {train_f1:.4f}")
            print(f"Val Loss: {val_metrics['loss']:.4f} | Val Acc: {val_metrics['accuracy']:.4f} | Val F1: {val_metrics['f1']:.4f}")
            print(f"Val Precision: {val_metrics['precision']:.4f} | Val Recall: {val_metrics['recall']:.4f}")
            print(f"Learning Rate: {scheduler.get_last_lr()[0]:.2e}")
            
            # Save best model
            if val_metrics['f1'] > best_val_f1:
                best_val_f1 = val_metrics['f1']
                patience_counter = 0
                self.save_model('models/best_model.pt', val_metrics, epoch)
                print(f"✓ New best model saved! (F1: {best_val_f1:.4f})")
            else:
                patience_counter += 1
                print(f"⚠ No improvement ({patience_counter}/{self.patience})")
            
            # Early stopping
            if patience_counter >= self.patience:
                print(f"\n⚠ Early stopping triggered after {epoch+1} epochs")
                break
            
            print("-" * 60)
        
        return history
    
    def evaluate(self, data_loader, criterion=None):
        """Evaluate model"""
        self.model.eval()
        
        if criterion is None:
            criterion = nn.CrossEntropyLoss()
        
        all_preds, all_labels, all_probs = [], [], []
        total_loss = 0
        
        with torch.no_grad():
            for batch in data_loader:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                adjacency_matrix = batch['adjacency_matrix'].to(self.device)
                labels = batch['label'].to(self.device)
                numeric_features = batch.get('numeric_features', None)
                if numeric_features is not None:
                    numeric_features = numeric_features.to(self.device)
                
                logits = self.model(input_ids, attention_mask, adjacency_matrix, numeric_features)
                loss = criterion(logits, labels)
                total_loss += loss.item()
                
                probs = F.softmax(logits, dim=-1)
                preds = torch.argmax(logits, dim=-1)
                
                all_preds.extend(preds.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
                all_probs.extend(probs[:, 1].cpu().numpy())
        
        metrics = {
            'loss': total_loss / len(data_loader),
            'accuracy': accuracy_score(all_labels, all_preds),
            'precision': precision_score(all_labels, all_preds, zero_division=0),
            'recall': recall_score(all_labels, all_preds, zero_division=0),
            'f1': f1_score(all_labels, all_preds, zero_division=0),
            'auc': roc_auc_score(all_labels, all_probs) if len(set(all_labels)) > 1 else 0.0,
            'predictions': all_preds,
            'labels': all_labels,
            'probabilities': all_probs
        }
        
        return metrics
    
    def save_model(self, path, metrics=None, epoch=None):
        """Save model with metadata"""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'model_config': {
                'numeric_feature_dim': self.model.numeric_feature_dim,
                'num_gat_layers': len(self.model.gat_layers),
            },
            'training_config': {
                'learning_rate': self.learning_rate,
                'weight_decay': self.weight_decay,
                'batch_size': self.batch_size,
                'max_length': self.max_length
            },
            'metrics': metrics,
            'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        
        torch.save(checkpoint, path)
    
    def load_model(self, path):
        """Load model checkpoint"""
        checkpoint = torch.load(path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        print(f"✓ Model loaded from {path}")
        if 'metrics' in checkpoint and checkpoint['metrics']:
            print(f"  Best F1: {checkpoint['metrics'].get('f1', 'N/A'):.4f}")
        return checkpoint


def plot_training_history(history, save_path='results/training_curves.png'):
    """Plot and save training curves"""
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    
    # Loss
    axes[0, 0].plot(history['train_loss'], label='Train Loss', marker='o')
    axes[0, 0].plot(history['val_loss'], label='Val Loss', marker='s')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].set_title('Training and Validation Loss')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    
    # Accuracy
    axes[0, 1].plot(history['train_acc'], label='Train Acc', marker='o')
    axes[0, 1].plot(history['val_acc'], label='Val Acc', marker='s')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Accuracy')
    axes[0, 1].set_title('Training and Validation Accuracy')
    axes[0, 1].legend()
    axes[0, 1].grid(True, alpha=0.3)
    
    # F1 Score
    axes[1, 0].plot(history['train_f1'], label='Train F1', marker='o')
    axes[1, 0].plot(history['val_f1'], label='Val F1', marker='s')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('F1 Score')
    axes[1, 0].set_title('Training and Validation F1 Score')
    axes[1, 0].legend()
    axes[1, 0].grid(True, alpha=0.3)
    
    # Learning Rate
    axes[1, 1].plot(history['learning_rates'], marker='o', color='purple')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('Learning Rate')
    axes[1, 1].set_title('Learning Rate Schedule')
    axes[1, 1].set_yscale('log')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"✓ Training curves saved to {save_path}")
    plt.close()


def plot_confusion_matrix(labels, predictions, save_path='results/confusion_matrix.png'):
    """Plot and save confusion matrix"""
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    cm = confusion_matrix(labels, predictions)
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['Non-Vulnerable', 'Vulnerable'],
                yticklabels=['Non-Vulnerable', 'Vulnerable'])
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix - Test Set')
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"✓ Confusion matrix saved to {save_path}")
    plt.close()


def save_results(history, test_metrics, save_path='results/training_results.json'):
    """Save all results to JSON"""
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    results = {
        'training_history': {
            'train_loss': [float(x) for x in history['train_loss']],
            'train_accuracy': [float(x) for x in history['train_acc']],
            'train_f1': [float(x) for x in history['train_f1']],
            'val_loss': [float(x) for x in history['val_loss']],
            'val_accuracy': [float(x) for x in history['val_acc']],
            'val_f1': [float(x) for x in history['val_f1']],
        },
        'test_metrics': {
            'accuracy': float(test_metrics['accuracy']),
            'precision': float(test_metrics['precision']),
            'recall': float(test_metrics['recall']),
            'f1_score': float(test_metrics['f1']),
            'auc_roc': float(test_metrics['auc'])
        },
        'best_epoch': {
            'best_val_f1': float(max(history['val_f1'])),
            'epoch_number': int(np.argmax(history['val_f1']) + 1)
        },
        'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }
    
    with open(save_path, 'w') as f:
        json.dump(results, f, indent=4)
    
    print(f"✓ Results saved to {save_path}")


def visualize_sample_graphs(train_df, num_samples=3):
    """Visualize AST and DFG for sample code before training"""
    print("\n" + "="*60)
    print("VISUALIZING SAMPLE AST+DFG GRAPHS")
    print("="*60)
    
    from __init__ import GraphBuilder
    
    os.makedirs('visualizations', exist_ok=True)
    
    builder = GraphBuilder()
    
    # Get samples from both classes
    vuln_samples = train_df[train_df['label'] == 1].head(num_samples)
    non_vuln_samples = train_df[train_df['label'] == 0].head(num_samples)
    
    print(f"\n✓ Visualizing {num_samples} vulnerable and {num_samples} non-vulnerable samples")
    
    for idx, (_, row) in enumerate(vuln_samples.iterrows()):
        code = row['code']
        print(f"\nProcessing vulnerable sample {idx+1}...")
        try:
            builder.visualize_combined(
                code, 
                save_path=f'visualizations/vulnerable_sample_{idx+1}.png',
                max_ast_depth=3
            )
        except Exception as e:
            print(f"  Warning: Could not visualize sample {idx+1}: {e}")
    
    for idx, (_, row) in enumerate(non_vuln_samples.iterrows()):
        code = row['code']
        print(f"\nProcessing non-vulnerable sample {idx+1}...")
        try:
            builder.visualize_combined(
                code, 
                save_path=f'visualizations/non_vulnerable_sample_{idx+1}.png',
                max_ast_depth=3
            )
        except Exception as e:
            print(f"  Warning: Could not visualize sample {idx+1}: {e}")
    
    print(f"\n✓ Visualizations saved to visualizations/ directory")
    print("="*60)


def visualize_adjacency_matrix_sample(train_loader):
    """Visualize a sample adjacency matrix to show graph structure"""
    print("\n" + "="*60)
    print("VISUALIZING SAMPLE ADJACENCY MATRIX")
    print("="*60)
    
    os.makedirs('visualizations', exist_ok=True)
    
    # Get one batch
    for batch in train_loader:
        adj_matrix = batch['adjacency_matrix'][0].cpu().numpy()  # First sample
        
        # Only visualize first 50x50 for clarity
        adj_viz = adj_matrix[:50, :50]
        
        plt.figure(figsize=(12, 10))
        
        # Create custom colormap: 0=white, 1=blue (AST), 2=red (DFG)
        from matplotlib.colors import ListedColormap
        colors = ['white', '#4A90E2', '#E74C3C']  # White, Blue, Red
        cmap = ListedColormap(colors)
        
        plt.imshow(adj_viz, cmap=cmap, interpolation='nearest', vmin=0, vmax=2)
        plt.colorbar(ticks=[0, 1, 2], label='Edge Type', 
                     format=plt.FuncFormatter(lambda x, p: ['None', 'AST', 'DFG'][int(x)]))
        
        plt.title('Adjacency Matrix Visualization (50x50 subset)\nBlue: AST edges | Red: DFG edges', 
                  fontsize=14, fontweight='bold', pad=20)
        plt.xlabel('Node Index')
        plt.ylabel('Node Index')
        
        # Add grid
        plt.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
        
        plt.tight_layout()
        plt.savefig('visualizations/adjacency_matrix_sample.png', dpi=300, bbox_inches='tight')
        print("✓ Adjacency matrix visualization saved to visualizations/adjacency_matrix_sample.png")
        plt.close()
        
        # Print statistics
        num_ast_edges = np.sum(adj_viz == 1)
        num_dfg_edges = np.sum(adj_viz == 2)
        print(f"\nAdjacency Matrix Statistics (50x50 region):")
        print(f"  AST edges (parent-child): {num_ast_edges}")
        print(f"  DFG edges (data flow): {num_dfg_edges}")
        print(f"  Total edges: {num_ast_edges + num_dfg_edges}")
        print(f"  Sparsity: {((2500 - num_ast_edges - num_dfg_edges) / 2500 * 100):.2f}%")
        
        break  # Only process one sample
    
    print("="*60)


def main():
    """Main training pipeline"""
    
    print("="*60)
    print("JAVASCRIPT VULNERABILITY DETECTION")
    print("GraphCodeBERT with AST+DFG")
    print("="*60)
    
    # Create directories
    os.makedirs('models', exist_ok=True)
    os.makedirs('results', exist_ok=True)
    os.makedirs('visualizations', exist_ok=True)
    
    # Load preprocessed data
    print("\n" + "="*60)
    print("LOADING PREPROCESSED DATA")
    print("="*60)
    
    try:
        train_df = pd.read_csv('processed_data/train.csv')
        val_df = pd.read_csv('processed_data/val.csv')
        test_df = pd.read_csv('processed_data/test.csv')
        
        train_features = np.load('processed_data/train_features.npy')
        val_features = np.load('processed_data/val_features.npy')
        test_features = np.load('processed_data/test_features.npy')
        
        print(f"✓ Train samples: {len(train_df)}")
        print(f"✓ Val samples: {len(val_df)}")
        print(f"✓ Test samples: {len(test_df)}")
        print(f"✓ Feature dimension: {train_features.shape[1]}")
        
    except FileNotFoundError as e:
        print(f"\n❌ ERROR: Preprocessed data not found!")
        print(f"Please run 'python data_processing.py' first to generate processed data.")
        print(f"Missing file: {e.filename}")
        return
    except Exception as e:
        print(f"\n❌ ERROR loading data: {e}")
        return
    
    # Visualize sample graphs BEFORE training
    try:
        visualize_sample_graphs(train_df, num_samples=3)
    except Exception as e:
        print(f"⚠ Warning: Could not generate visualizations: {e}")
    
    # Detect hardware and optimize
    if torch.cuda.is_available():
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
        if gpu_mem < 6:
            batch_size = 4
        elif gpu_mem < 8:
            batch_size = 8
        else:
            batch_size = 16
        print(f"\n✓ GPU detected. Using batch_size={batch_size}")
    else:
        batch_size = 4
        print(f"\n✓ CPU mode. Using batch_size={batch_size}")
    
    # Initialize detector
    print("\n" + "="*60)
    print("INITIALIZING MODEL")
    print("="*60)
    
    try:
        detector = VulnerabilityDetector(
            model_name='microsoft/graphcodebert-base',
            numeric_feature_dim=train_features.shape[1],
            learning_rate=2e-5,
            weight_decay=0.01,
            num_epochs=15,
            batch_size=batch_size,
            max_length=512,
            patience=5
        )
    except Exception as e:
        print(f"\n❌ ERROR initializing model: {e}")
        print("Make sure transformers and torch are properly installed.")
        return
    
    # Prepare data loaders
    try:
        train_loader, val_loader, test_loader = detector.prepare_data_loaders(
            train_data=(train_df, train_df['label'], train_features),
            val_data=(val_df, val_df['label'], val_features),
            test_data=(test_df, test_df['label'], test_features)
        )
    except KeyError as e:
        print(f"\n❌ ERROR: Missing column {e} in dataframe")
        print(f"Available columns: {train_df.columns.tolist()}")
        return
    except Exception as e:
        print(f"\n❌ ERROR preparing data loaders: {e}")
        return
    
    # Visualize adjacency matrix structure
    try:
        visualize_adjacency_matrix_sample(train_loader)
    except Exception as e:
        print(f"⚠ Warning: Could not visualize adjacency matrix: {e}")
    
    # Train model
    try:
        history = detector.train(train_loader, val_loader)
    except Exception as e:
        print(f"\n❌ ERROR during training: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # Plot training curves
    try:
        plot_training_history(history)
    except Exception as e:
        print(f"⚠ Warning: Could not plot training curves: {e}")
    
    # Load best model and evaluate on test set
    print("\n" + "="*60)
    print("FINAL EVALUATION ON TEST SET")
    print("="*60)
    
    try:
        detector.load_model('models/best_model.pt')
        test_metrics = detector.evaluate(test_loader)
        
        print(f"\n{'='*60}")
        print("TEST SET RESULTS")
        print(f"{'='*60}")
        print(f"Accuracy:  {test_metrics['accuracy']:.4f}")
        print(f"Precision: {test_metrics['precision']:.4f}")
        print(f"Recall:    {test_metrics['recall']:.4f}")
        print(f"F1 Score:  {test_metrics['f1']:.4f}")
        print(f"AUC-ROC:   {test_metrics['auc']:.4f}")
        print(f"{'='*60}")
        
        # Plot confusion matrix
        plot_confusion_matrix(test_metrics['labels'], test_metrics['predictions'])
        
        # Save all results
        save_results(history, test_metrics)
        
    except FileNotFoundError:
        print(f"\n⚠ Warning: Best model not found. Training may have failed.")
    except Exception as e:
        print(f"\n❌ ERROR during evaluation: {e}")
        import traceback
        traceback.print_exc()
        return
    
    print("\n" + "="*60)
    print("TRAINING COMPLETE!")
    print("="*60)
    print("✓ Model saved: models/best_model.pt")
    print("✓ Training curves: results/training_curves.png")
    print("✓ Confusion matrix: results/confusion_matrix.png")
    print("✓ Results JSON: results/training_results.json")
    print("✓ Sample visualizations: visualizations/")
    print("="*60)


if __name__ == "__main__":
    main()