train_models.py

"""
Comprehensive Training Script for Multimodal Parkinson's Detection System
Trains voice, handwriting, and fusion models with proper validation
"""

import os
import sys
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, List, Tuple, Optional
import json
import joblib
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

# Add src to path
sys.path.append('src')

from models.voice_model import VoiceModel
from models.handwriting_model import HandwritingCNN, HandwritingDataProcessor
from models.fusion_model import MultimodalFusion
from explainability.grad_cam import HandwritingExplainer
from utils.pdf_report import ClinicalReportGenerator


class HandwritingDataset(Dataset):
    """Dataset class for handwriting images"""
    
    def __init__(self, image_paths: List[str], labels: List[int], processor, transform=None):
        self.image_paths = image_paths
        self.labels = labels
        self.processor = processor
        self.transform = transform
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        image_path = self.image_paths[idx]
        label = self.labels[idx]
        
        # Process image
        image = self.processor.preprocess_image(image_path)
        
        if self.transform:
            image = self.transform(image)
        
        return image, torch.tensor(label, dtype=torch.long)


def prepare_handwriting_data(data_dir: str) -> Tuple[List[str], List[int]]:
    """
    Prepare handwriting dataset
    
    Args:
        data_dir: Directory containing handwriting images
        
    Returns:
        Tuple of (image_paths, labels)
    """
    image_paths = []
    labels = []
    
    # Healthy samples
    healthy_dir = os.path.join(data_dir, 'Healthy')
    if os.path.exists(healthy_dir):
        for filename in os.listdir(healthy_dir):
            if filename.endswith(('.png', '.jpg', '.jpeg')):
                image_paths.append(os.path.join(healthy_dir, filename))
                labels.append(0)  # Healthy = 0
    
    # Parkinson samples
    parkinson_dir = os.path.join(data_dir, 'Parkinson')
    if os.path.exists(parkinson_dir):
        for filename in os.listdir(parkinson_dir):
            if filename.endswith(('.png', '.jpg', '.jpeg')):
                image_paths.append(os.path.join(parkinson_dir, filename))
                labels.append(1)  # Parkinson = 1
    
    return image_paths, labels


def train_handwriting_model(image_paths: List[str], 
                           labels: List[int],
                           config: Dict,
                           device: torch.device) -> Tuple[HandwritingCNN, Dict]:
    """
    Train handwriting CNN model
    
    Args:
        image_paths: List of image paths
        labels: List of labels
        config: Training configuration
        device: Device to train on
        
    Returns:
        Tuple of (trained_model, training_results)
    """
    print("Training Handwriting CNN Model...")
    
    # Create processor
    processor = HandwritingDataProcessor(image_size=config['image_size'])
    
    # Split data
    train_paths, val_paths, train_labels, val_labels = train_test_split(
        image_paths, labels, test_size=0.2, random_state=42, stratify=labels
    )
    
    # Create datasets
    train_dataset = HandwritingDataset(train_paths, train_labels, processor)
    val_dataset = HandwritingDataset(val_paths, val_labels, processor)
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset, 
        batch_size=config['batch_size'], 
        shuffle=True,
        num_workers=2
    )
    val_loader = DataLoader(
        val_dataset, 
        batch_size=config['batch_size'], 
        shuffle=False,
        num_workers=2
    )
    
    # Create model
    model = HandwritingCNN(
        model_name=config['model_name'],
        num_classes=2,
        dropout_rate=config['dropout_rate'],
        pretrained=config['pretrained']
    ).to(device)
    
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(
        model.parameters(), 
        lr=config['learning_rate'],
        weight_decay=config['weight_decay']
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', patience=5, factor=0.5
    )
    
    # Training loop
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []
    
    best_val_acc = 0
    best_model_state = None
    
    for epoch in range(config['epochs']):
        # Training phase
        model.train()
        train_loss = 0
        train_correct = 0
        train_total = 0
        
        for images, labels in tqdm(train_loader, desc=f'Epoch {epoch+1}/{config["epochs"]}'):
            images, labels = images.to(device), labels.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs['logits'], labels)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = torch.max(outputs['logits'], 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()
        
        train_loss /= len(train_loader)
        train_acc = train_correct / train_total
        
        # Validation phase
        model.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs['logits'], labels)
                
                val_loss += loss.item()
                _, predicted = torch.max(outputs['logits'], 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()
        
        val_loss /= len(val_loader)
        val_acc = val_correct / val_total
        
        # Update learning rate
        scheduler.step(val_loss)
        
        # Store metrics
        train_losses.append(train_loss)
        val_losses.append(val_loss)
        train_accuracies.append(train_acc)
        val_accuracies.append(val_acc)
        
        print(f'Epoch {epoch+1}: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, '
              f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')
        
        # Save best model
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_model_state = model.state_dict().copy()
    
    # Load best model
    model.load_state_dict(best_model_state)
    
    # Final evaluation
    model.eval()
    all_predictions = []
    all_labels = []
    
    with torch.no_grad():
        for images, labels in val_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs['logits'], 1)
            
            all_predictions.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # Calculate metrics
    accuracy = accuracy_score(all_labels, all_predictions)
    report = classification_report(all_labels, all_predictions, output_dict=True)
    cm = confusion_matrix(all_labels, all_predictions)
    
    results = {
        'accuracy': accuracy,
        'classification_report': report,
        'confusion_matrix': cm,
        'train_losses': train_losses,
        'val_losses': val_losses,
        'train_accuracies': train_accuracies,
        'val_accuracies': val_accuracies,
        'best_val_accuracy': best_val_acc
    }
    
    return model, results


def train_voice_model(csv_path: str, config: Dict) -> Tuple[VoiceModel, Dict]:
    """
    Train voice model
    
    Args:
        csv_path: Path to voice data CSV
        config: Training configuration
        
    Returns:
        Tuple of (trained_model, training_results)
    """
    print("Training Voice Model...")
    
    # Create model
    model = VoiceModel(
        model_type=config['model_type'],
        use_calibration=config['use_calibration'],
        feature_engineering=config['feature_engineering']
    )
    
    # Load data
    X, y = model.load_data(csv_path)
    
    # Train model
    results = model.train(X, y, test_size=0.2)
    
    return model, results


def train_fusion_model(voice_model: VoiceModel,
                      handwriting_model: HandwritingCNN,
                      voice_data: pd.DataFrame,
                      handwriting_data: Tuple[List[str], List[int]],
                      config: Dict,
                      device: torch.device) -> Tuple[MultimodalFusion, Dict]:
    """
    Train fusion model
    
    Args:
        voice_model: Trained voice model
        handwriting_model: Trained handwriting model
        voice_data: Voice data
        handwriting_data: Handwriting data (paths, labels)
        config: Training configuration
        device: Device to use
        
    Returns:
        Tuple of (trained_fusion_model, training_results)
    """
    print("Training Fusion Model...")
    
    # Create processor
    processor = HandwritingDataProcessor()
    
    # Generate predictions for training fusion model
    voice_predictions = []
    handwriting_predictions = []
    voice_uncertainties = []
    handwriting_uncertainties = []
    labels = []
    
    # Process voice data
    voice_features = voice_data.drop('status', axis=1)
    voice_labels = voice_data['status']
    
    if voice_model.feature_engineering:
        voice_features = voice_model.engineer_features(voice_features)
    
    voice_features_scaled = voice_model.scaler.transform(voice_features)
    
    for i in range(len(voice_features_scaled)):
        # Voice prediction
        voice_pred_proba = voice_model.predict_proba(voice_features_scaled[i:i+1])[0]
        voice_pred = {
            'probabilities': voice_pred_proba,
            'prediction': np.argmax(voice_pred_proba),
            'embeddings': np.zeros(128)  # Placeholder
        }
        voice_predictions.append(voice_pred)
        
        # Voice uncertainty (placeholder)
        voice_uncertainty = {
            'entropy': [0.5],
            'confidence': [0.5]
        }
        voice_uncertainties.append(voice_uncertainty)
    
    # Process handwriting data (sample subset for efficiency)
    image_paths, handwriting_labels = handwriting_data
    n_samples = min(len(image_paths), len(voice_predictions))
    indices = np.random.choice(len(image_paths), n_samples, replace=False)
    
    handwriting_model.eval()
    with torch.no_grad():
        for idx in indices:
            image_path = image_paths[idx]
            label = handwriting_labels[idx]
            
            # Process image
            image = processor.preprocess_image(image_path).unsqueeze(0).to(device)
            
            # Handwriting prediction
            handwriting_output = handwriting_model(image)
            handwriting_uncertainty = handwriting_model.predict_with_uncertainty(image)
            
            handwriting_pred = {
                'probabilities': handwriting_output['probabilities'][0].cpu().numpy(),
                'prediction': torch.argmax(handwriting_output['logits'][0]).item(),
                'embeddings': handwriting_output['embeddings'][0].cpu().numpy()
            }
            handwriting_predictions.append(handwriting_pred)
            
            handwriting_uncertainty_dict = {
                'entropy': handwriting_uncertainty['entropy'].cpu().numpy(),
                'confidence': handwriting_uncertainty['confidence'].cpu().numpy()
            }
            handwriting_uncertainties.append(handwriting_uncertainty_dict)
            
            labels.append(label)
    
    # Align voice and handwriting data
    voice_predictions = voice_predictions[:len(handwriting_predictions)]
    voice_uncertainties = voice_uncertainties[:len(handwriting_uncertainties)]
    
    # Create fusion model
    fusion_model = MultimodalFusion(
        fusion_method=config['fusion_method'],
        use_uncertainty=config['use_uncertainty'],
        use_embeddings=config['use_embeddings']
    )
    
    # Train fusion model
    results = fusion_model.train(
        voice_predictions,
        handwriting_predictions,
        np.array(labels),
        voice_uncertainties,
        handwriting_uncertainties
    )
    
    return fusion_model, results


def save_models(voice_model: VoiceModel,
                handwriting_model: HandwritingCNN,
                fusion_model: MultimodalFusion,
                results: Dict,
                save_dir: str = 'models'):
    """
    Save trained models and results
    
    Args:
        voice_model: Trained voice model
        handwriting_model: Trained handwriting model
        fusion_model: Trained fusion model
        results: Training results
        save_dir: Directory to save models
    """
    os.makedirs(save_dir, exist_ok=True)
    
    # Save voice model
    voice_model.save_model(os.path.join(save_dir, 'voice_model.pkl'))
    
    # Save handwriting model
    torch.save(handwriting_model.state_dict(), os.path.join(save_dir, 'handwriting_model.pth'))
    
    # Save fusion model
    fusion_model.save_model(os.path.join(save_dir, 'fusion_model.pkl'))
    
    # Save results
    with open(os.path.join(save_dir, 'training_results.json'), 'w') as f:
        # Convert numpy arrays to lists for JSON serialization
        json_results = {}
        for key, value in results.items():
            if isinstance(value, np.ndarray):
                json_results[key] = value.tolist()
            elif isinstance(value, dict):
                json_results[key] = {}
                for k, v in value.items():
                    if isinstance(v, np.ndarray):
                        json_results[key][k] = v.tolist()
                    else:
                        json_results[key][k] = v
            else:
                json_results[key] = value
        
        json.dump(json_results, f, indent=2)
    
    print(f"Models and results saved to {save_dir}/")


def create_training_plots(results: Dict, save_dir: str = 'models'):
    """
    Create training visualization plots
    
    Args:
        results: Training results
        save_dir: Directory to save plots
    """
    # Handwriting model training curves
    if 'handwriting_results' in results:
        hw_results = results['handwriting_results']
        
        fig, ((ax1, ax2)) = plt.subplots(1, 2, figsize=(12, 5))
        
        # Loss curves
        ax1.plot(hw_results['train_losses'], label='Train Loss')
        ax1.plot(hw_results['val_losses'], label='Validation Loss')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.set_title('Handwriting Model Training Loss')
        ax1.legend()
        ax1.grid(True)
        
        # Accuracy curves
        ax2.plot(hw_results['train_accuracies'], label='Train Accuracy')
        ax2.plot(hw_results['val_accuracies'], label='Validation Accuracy')
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy')
        ax2.set_title('Handwriting Model Training Accuracy')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        plt.savefig(os.path.join(save_dir, 'handwriting_training_curves.png'), dpi=300, bbox_inches='tight')
        plt.close()
        
        # Confusion matrix
        cm = np.array(hw_results['confusion_matrix'])
        plt.figure(figsize=(8, 6))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=['Healthy', 'Parkinson\'s'],
                   yticklabels=['Healthy', 'Parkinson\'s'])
        plt.title('Handwriting Model Confusion Matrix')
        plt.ylabel('True Label')
        plt.xlabel('Predicted Label')
        plt.savefig(os.path.join(save_dir, 'handwriting_confusion_matrix.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    # Voice model results
    if 'voice_results' in results:
        voice_results = results['voice_results']
        
        # Feature importance
        if hasattr(voice_results, 'feature_importance'):
            importance_df = voice_results['feature_importance'].head(15)
            plt.figure(figsize=(10, 8))
            sns.barplot(data=importance_df, x='importance', y='feature')
            plt.title('Voice Model Feature Importance')
            plt.xlabel('Importance')
            plt.tight_layout()
            plt.savefig(os.path.join(save_dir, 'voice_feature_importance.png'), dpi=300, bbox_inches='tight')
            plt.close()
    
    # Fusion model results
    if 'fusion_results' in results:
        fusion_results = results['fusion_results']
        
        # Feature importance
        if 'feature_importance' in fusion_results:
            importance_df = fusion_results['feature_importance'].head(15)
            plt.figure(figsize=(12, 8))
            sns.barplot(data=importance_df, x='importance', y='feature')
            plt.title('Fusion Model Feature Importance')
            plt.xlabel('Importance')
            plt.xticks(rotation=45)
            plt.tight_layout()
            plt.savefig(os.path.join(save_dir, 'fusion_feature_importance.png'), dpi=300, bbox_inches='tight')
            plt.close()
    
    print(f"Training plots saved to {save_dir}/")


def main():
    """Main training function"""
    
    # Configuration
    config = {
        'handwriting': {
            'model_name': 'resnet18',
            'image_size': 224,
            'batch_size': 32,
            'epochs': 20,
            'learning_rate': 0.001,
            'weight_decay': 1e-4,
            'dropout_rate': 0.3,
            'pretrained': True
        },
        'voice': {
            'model_type': 'xgboost',
            'use_calibration': True,
            'feature_engineering': True
        },
        'fusion': {
            'fusion_method': 'xgboost',
            'use_uncertainty': True,
            'use_embeddings': True
        }
    }
    
    # Device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # Paths
    voice_csv_path = 'parkinsons1.csv'
    handwriting_data_dir = 'Parkinsons Hand Written Samples/Dataset/Dataset'
    
    # Check if data exists
    if not os.path.exists(voice_csv_path):
        print(f"Voice data not found: {voice_csv_path}")
        return
    
    if not os.path.exists(handwriting_data_dir):
        print(f"Handwriting data not found: {handwriting_data_dir}")
        return
    
    # Prepare data
    print("Preparing data...")
    handwriting_paths, handwriting_labels = prepare_handwriting_data(handwriting_data_dir)
    print(f"Found {len(handwriting_paths)} handwriting images")
    
    # Load voice data
    voice_data = pd.read_csv(voice_csv_path)
    if 'name' in voice_data.columns:
        voice_data = voice_data.drop('name', axis=1)
    print(f"Found {len(voice_data)} voice samples")
    
    # Train models
    results = {}
    
    # Train voice model
    voice_model, voice_results = train_voice_model(voice_csv_path, config['voice'])
    results['voice_results'] = voice_results
    print(f"Voice model accuracy: {voice_results['test_accuracy']:.4f}")
    
    # Train handwriting model
    handwriting_model, handwriting_results = train_handwriting_model(
        handwriting_paths, handwriting_labels, config['handwriting'], device
    )
    results['handwriting_results'] = handwriting_results
    print(f"Handwriting model accuracy: {handwriting_results['accuracy']:.4f}")
    
    # Train fusion model
    fusion_model, fusion_results = train_fusion_model(
        voice_model, handwriting_model, voice_data, 
        (handwriting_paths, handwriting_labels), config['fusion'], device
    )
    results['fusion_results'] = fusion_results
    print(f"Fusion model accuracy: {fusion_results['train_accuracy']:.4f}")
    
    # Save models and results
    save_models(voice_model, handwriting_model, fusion_model, results)
    
    # Create training plots
    create_training_plots(results)
    
    # Summary
    print("\n" + "="*50)
    print("TRAINING SUMMARY")
    print("="*50)
    print(f"Voice Model Accuracy: {voice_results['test_accuracy']:.4f}")
    print(f"Handwriting Model Accuracy: {handwriting_results['accuracy']:.4f}")
    print(f"Fusion Model Accuracy: {fusion_results['train_accuracy']:.4f}")
    print("\nAll models saved to 'models/' directory")
    print("Training plots saved to 'models/' directory")
    print("Ready to run the Streamlit app!")


if __name__ == "__main__":
    main()