Trixy3/trainer/utils.py

"""
Utility classes for training progress monitoring, logging, and validation.

This module provides comprehensive utilities for tracking training progress,
professional logging, performance monitoring, and validation metrics.
"""

import os
import time
import json
import logging
import threading
from pathlib import Path
from typing import Dict, List, Any, Optional, Callable, Union
from dataclasses import dataclass, field
from datetime import datetime, timedelta
import sys

import torch
import torch.nn as nn
import numpy as np
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support, confusion_matrix,
    classification_report, roc_auc_score, average_precision_score
)
try:
    import matplotlib.pyplot as plt
    import seaborn as sns
    PLOTTING_AVAILABLE = True
except ImportError:
    PLOTTING_AVAILABLE = False

try:
    from tensorboard import SummaryWriter
    TENSORBOARD_AVAILABLE = True
except ImportError:
    TENSORBOARD_AVAILABLE = False
    SummaryWriter = None

try:
    import wandb
    WANDB_AVAILABLE = True
except ImportError:
    WANDB_AVAILABLE = False
    wandb = None


def make_json_serializable(obj: Any) -> Any:
    """
    Convert numpy types and other non-JSON-serializable objects to JSON-safe types.
    
    Args:
        obj: Object to convert
        
    Returns:
        JSON-serializable object
    """
    # Handle numpy integer types
    if isinstance(obj, (np.integer,)):
        return int(obj)
    # Handle numpy floating point types
    elif isinstance(obj, (np.floating,)):
        return float(obj)
    # Handle numpy boolean types
    elif isinstance(obj, np.bool_):
        return bool(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, torch.Tensor):
        return obj.detach().cpu().numpy().tolist()
    elif isinstance(obj, dict):
        return {key: make_json_serializable(value) for key, value in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [make_json_serializable(item) for item in obj]
    elif hasattr(obj, '__dict__'):
        # Handle dataclass or custom objects
        if hasattr(obj, 'to_dict'):
            return make_json_serializable(obj.to_dict())
        else:
            return make_json_serializable(obj.__dict__)
    else:
        return obj


class TrainerLogger:
    """
    Professional logging system for ML trainers.
    
    Provides structured logging with file output, console output,
    and optional integration with TensorBoard and Weights & Biases.
    """
    
    def __init__(self, name: str, log_dir: str = "./logs", 
                 log_level: str = "INFO", 
                 use_tensorboard: bool = True,
                 use_wandb: bool = False,
                 wandb_project: Optional[str] = None):
        """
        Initialize trainer logger.
        
        Args:
            name: Logger name
            log_dir: Directory for log files
            log_level: Logging level
            use_tensorboard: Whether to use TensorBoard logging
            use_wandb: Whether to use Weights & Biases logging
            wandb_project: W&B project name
        """
        self.name = name
        self.log_dir = Path(log_dir)
        self.log_dir.mkdir(parents=True, exist_ok=True)
        
        # Setup basic logger
        self.logger = self._setup_logger(log_level)
        
        # Setup TensorBoard
        self.tb_writer = None
        if use_tensorboard and TENSORBOARD_AVAILABLE:
            tb_log_dir = self.log_dir / "tensorboard" / name
            tb_log_dir.mkdir(parents=True, exist_ok=True)
            self.tb_writer = SummaryWriter(str(tb_log_dir))
            self.logger.info("TensorBoard logging enabled")
        elif use_tensorboard:
            self.logger.warning("TensorBoard not available - install with: pip install tensorboard")
        
        # Setup Weights & Biases
        self.use_wandb = False
        if use_wandb and WANDB_AVAILABLE:
            if wandb_project:
                wandb.init(project=wandb_project, name=name)
                self.use_wandb = True
                self.logger.info("Weights & Biases logging enabled")
            else:
                self.logger.warning("W&B project name required for W&B logging")
        elif use_wandb:
            self.logger.warning("Weights & Biases not available - install with: pip install wandb")
        
        # Training metrics storage
        self.metrics_history = []
        self.current_step = 0
    
    def _setup_logger(self, log_level: str) -> logging.Logger:
        """Setup the basic logger with file and console handlers."""
        logger = logging.getLogger(self.name)
        logger.setLevel(getattr(logging, log_level.upper()))
        
        # Clear existing handlers
        logger.handlers.clear()
        
        # Create formatters
        detailed_formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        simple_formatter = logging.Formatter(
            '%(asctime)s - %(levelname)s - %(message)s'
        )
        
        # File handler (detailed)
        log_file = self.log_dir / f"{self.name}.log"
        file_handler = logging.FileHandler(log_file)
        file_handler.setLevel(logging.DEBUG)
        file_handler.setFormatter(detailed_formatter)
        logger.addHandler(file_handler)
        
        # Console handler (simple)
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(logging.INFO)
        console_handler.setFormatter(simple_formatter)
        logger.addHandler(console_handler)
        
        return logger
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None):
        """Log training metrics to all enabled loggers."""
        if step is None:
            step = self.current_step
            self.current_step += 1
        
        # Store metrics
        metric_entry = {"step": step, "timestamp": datetime.now().isoformat()}
        metric_entry.update(metrics)
        self.metrics_history.append(metric_entry)
        
        # Log to basic logger
        metric_str = ", ".join([f"{k}: {v:.6f}" for k, v in metrics.items()])
        self.logger.info(f"Step {step} - {metric_str}")
        
        # Log to TensorBoard
        if self.tb_writer:
            for key, value in metrics.items():
                self.tb_writer.add_scalar(key, value, step)
            self.tb_writer.flush()
        
        # Log to Weights & Biases
        if self.use_wandb:
            wandb.log(metrics, step=step)
    
    def log_model_graph(self, model: nn.Module, input_sample: torch.Tensor):
        """Log model graph to TensorBoard."""
        if self.tb_writer:
            try:
                self.tb_writer.add_graph(model, input_sample)
                self.logger.info("Model graph logged to TensorBoard")
            except Exception as e:
                self.logger.warning(f"Failed to log model graph: {str(e)}")
    
    def log_hyperparameters(self, hparams: Dict[str, Any]):
        """Log hyperparameters."""
        # Log to basic logger
        self.logger.info("Hyperparameters:")
        for key, value in hparams.items():
            self.logger.info(f"  {key}: {value}")
        
        # Log to TensorBoard
        if self.tb_writer:
            # Convert non-scalar values to strings
            tb_hparams = {}
            for key, value in hparams.items():
                if isinstance(value, (int, float, str, bool)):
                    tb_hparams[key] = value
                else:
                    tb_hparams[key] = str(value)
            self.tb_writer.add_hparams(tb_hparams, {})
        
        # Log to Weights & Biases
        if self.use_wandb:
            wandb.config.update(hparams)
    
    def log_confusion_matrix(self, y_true: np.ndarray, y_pred: np.ndarray, 
                           class_names: List[str], step: int):
        """Log confusion matrix visualization."""
        cm = confusion_matrix(y_true, y_pred)
        
        if not PLOTTING_AVAILABLE:
            self.logger.warning("Matplotlib not available, skipping confusion matrix visualization")
            return
        
        # Create matplotlib figure
        fig, ax = plt.subplots(figsize=(10, 8))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=class_names, yticklabels=class_names, ax=ax)
        ax.set_xlabel('Predicted')
        ax.set_ylabel('Actual')
        ax.set_title('Confusion Matrix')
        
        # Log to TensorBoard
        if self.tb_writer:
            self.tb_writer.add_figure('confusion_matrix', fig, step)
        
        # Log to Weights & Biases
        if self.use_wandb:
            wandb.log({"confusion_matrix": wandb.Image(fig)}, step=step)
        
        plt.close(fig)
    
    def save_metrics_history(self, filepath: Optional[str] = None):
        """Save metrics history to JSON file."""
        if filepath is None:
            filepath = self.log_dir / f"{self.name}_metrics.json"
        
        with open(filepath, 'w') as f:
            json.dump(self.metrics_history, f, indent=2)
        
        self.logger.info(f"Metrics history saved to: {filepath}")
    
    def close(self):
        """Close all logging handlers."""
        if self.tb_writer:
            self.tb_writer.close()
        
        if self.use_wandb:
            wandb.finish()
        
        # Save metrics history
        self.save_metrics_history()


class ProgressMonitor:
    """
    Real-time training progress monitor with ETA estimation.
    
    Provides progress tracking, ETA calculation, and performance monitoring
    during training with thread-safe updates.
    """
    
    def __init__(self, total_epochs: int, steps_per_epoch: int = None):
        """
        Initialize progress monitor.
        
        Args:
            total_epochs: Total number of training epochs
            steps_per_epoch: Number of steps per epoch (for step-level tracking)
        """
        self.total_epochs = total_epochs
        self.steps_per_epoch = steps_per_epoch
        self.total_steps = total_epochs * (steps_per_epoch or 1)
        
        # Progress tracking
        self.current_epoch = 0
        self.current_step = 0
        self.global_step = 0
        
        # Timing
        self.start_time = None
        self.epoch_start_time = None
        self.step_times = []
        self.epoch_times = []
        
        # Thread safety
        self.lock = threading.Lock()
        
        # Display settings
        self.update_interval = 1.0  # seconds
        self.last_update = 0
    
    def start_training(self):
        """Mark the start of training."""
        with self.lock:
            self.start_time = time.time()
            self.epoch_start_time = self.start_time
    
    def start_epoch(self, epoch: int):
        """Mark the start of an epoch."""
        with self.lock:
            self.current_epoch = epoch
            self.current_step = 0
            self.epoch_start_time = time.time()
    
    def update_step(self, step: int = None):
        """Update step progress."""
        with self.lock:
            if step is not None:
                self.current_step = step
            else:
                self.current_step += 1
            
            self.global_step += 1
            
            # Record step time
            current_time = time.time()
            if len(self.step_times) > 0:
                step_duration = current_time - self.step_times[-1]
            else:
                step_duration = current_time - self.epoch_start_time
            
            self.step_times.append(current_time)
            
            # Limit step times history
            if len(self.step_times) > 100:
                self.step_times = self.step_times[-100:]
    
    def end_epoch(self):
        """Mark the end of an epoch."""
        with self.lock:
            current_time = time.time()
            epoch_duration = current_time - self.epoch_start_time
            self.epoch_times.append(epoch_duration)
            
            # Limit epoch times history
            if len(self.epoch_times) > 50:
                self.epoch_times = self.epoch_times[-50:]
    
    def get_progress_info(self) -> Dict[str, Any]:
        """Get current progress information."""
        with self.lock:
            current_time = time.time()
            
            if self.start_time is None:
                return {"status": "not_started"}
            
            # Calculate progress
            epoch_progress = self.current_epoch / self.total_epochs
            if self.steps_per_epoch:
                step_progress = self.current_step / self.steps_per_epoch
                total_progress = (self.current_epoch + step_progress) / self.total_epochs
            else:
                total_progress = epoch_progress
            
            # Calculate timing
            elapsed_time = current_time - self.start_time
            
            # Estimate remaining time
            if total_progress > 0:
                eta = (elapsed_time / total_progress) * (1 - total_progress)
            else:
                eta = 0
            
            # Calculate average times
            avg_step_time = np.mean([
                self.step_times[i] - self.step_times[i-1] 
                for i in range(1, len(self.step_times))
            ]) if len(self.step_times) > 1 else 0
            
            avg_epoch_time = np.mean(self.epoch_times) if self.epoch_times else 0
            
            return {
                "status": "running",
                "current_epoch": self.current_epoch,
                "total_epochs": self.total_epochs,
                "current_step": self.current_step,
                "steps_per_epoch": self.steps_per_epoch,
                "global_step": self.global_step,
                "total_steps": self.total_steps,
                "epoch_progress": epoch_progress,
                "total_progress": total_progress,
                "elapsed_time": elapsed_time,
                "eta": eta,
                "avg_step_time": avg_step_time,
                "avg_epoch_time": avg_epoch_time,
                "elapsed_str": str(timedelta(seconds=int(elapsed_time))),
                "eta_str": str(timedelta(seconds=int(eta))),
            }
    
    def print_progress(self, metrics: Dict[str, float] = None, force: bool = False):
        """Print progress information to console."""
        current_time = time.time()
        
        if not force and (current_time - self.last_update) < self.update_interval:
            return
        
        self.last_update = current_time
        
        info = self.get_progress_info()
        
        if info["status"] != "running":
            return
        
        # Format progress bar
        bar_length = 40
        filled_length = int(bar_length * info["total_progress"])
        bar = "█" * filled_length + "░" * (bar_length - filled_length)
        
        # Format output
        output = f"\rEpoch {info['current_epoch']}/{info['total_epochs']} "
        output += f"[{bar}] {info['total_progress']*100:.1f}% "
        output += f"- {info['elapsed_str']} / {info['eta_str']} "
        
        if metrics:
            metric_str = " - ".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
            output += f"- {metric_str}"
        
        # Add step info if available
        if info["steps_per_epoch"]:
            output += f" - Step {info['current_step']}/{info['steps_per_epoch']}"
        
        print(output, end="", flush=True)


class ValidationMetrics:
    """
    Comprehensive validation and testing metrics calculator.
    
    Provides calculation of various performance metrics for classification
    and regression tasks with support for multi-class and multi-label scenarios.
    """
    
    def __init__(self, task_type: str = "classification", 
                 num_classes: int = None,
                 class_names: List[str] = None):
        """
        Initialize validation metrics calculator.
        
        Args:
            task_type: Type of task ('classification', 'regression', 'multilabel')
            num_classes: Number of classes for classification
            class_names: Names of classes
        """
        self.task_type = task_type
        self.num_classes = num_classes
        self.class_names = class_names or [f"Class_{i}" for i in range(num_classes or 2)]
    
    def calculate_classification_metrics(self, y_true: np.ndarray, 
                                       y_pred: np.ndarray, 
                                       y_prob: np.ndarray = None) -> Dict[str, Any]:
        """
        Calculate comprehensive classification metrics.
        
        Args:
            y_true: True labels
            y_pred: Predicted labels
            y_prob: Prediction probabilities (for AUC calculation)
            
        Returns:
            Dictionary containing various metrics
        """
        metrics = {}
        
        # Basic metrics
        metrics["accuracy"] = accuracy_score(y_true, y_pred)
        
        # Precision, recall, F1-score
        precision, recall, f1, support = precision_recall_fscore_support(
            y_true, y_pred, average=None, zero_division=0
        )
        
        # Per-class metrics
        for i, class_name in enumerate(self.class_names[:len(precision)]):
            metrics[f"precision_{class_name}"] = precision[i]
            metrics[f"recall_{class_name}"] = recall[i]
            metrics[f"f1_{class_name}"] = f1[i]
            metrics[f"support_{class_name}"] = support[i]
        
        # Averaged metrics
        metrics["precision_macro"] = np.mean(precision)
        metrics["recall_macro"] = np.mean(recall)
        metrics["f1_macro"] = np.mean(f1)
        
        # Weighted averages
        precision_weighted = np.average(precision, weights=support)
        recall_weighted = np.average(recall, weights=support)
        f1_weighted = np.average(f1, weights=support)
        
        metrics["precision_weighted"] = precision_weighted
        metrics["recall_weighted"] = recall_weighted
        metrics["f1_weighted"] = f1_weighted
        
        # Confusion matrix
        cm = confusion_matrix(y_true, y_pred)
        metrics["confusion_matrix"] = cm.tolist()
        
        # AUC metrics (if probabilities provided)
        if y_prob is not None:
            try:
                if len(np.unique(y_true)) == 2:
                    # Binary classification
                    metrics["auc_roc"] = roc_auc_score(y_true, y_prob[:, 1])
                    metrics["auc_pr"] = average_precision_score(y_true, y_prob[:, 1])
                else:
                    # Multi-class classification
                    metrics["auc_roc_ovr"] = roc_auc_score(
                        y_true, y_prob, multi_class='ovr', average='macro'
                    )
                    metrics["auc_roc_ovo"] = roc_auc_score(
                        y_true, y_prob, multi_class='ovo', average='macro'
                    )
            except Exception as e:
                # AUC calculation failed (e.g., only one class present)
                pass
        
        # Classification report
        report = classification_report(y_true, y_pred, 
                                     target_names=self.class_names[:len(np.unique(y_true))],
                                     output_dict=True, zero_division=0)
        metrics["classification_report"] = report
        
        return metrics
    
    def calculate_regression_metrics(self, y_true: np.ndarray, 
                                   y_pred: np.ndarray) -> Dict[str, float]:
        """
        Calculate regression metrics.
        
        Args:
            y_true: True values
            y_pred: Predicted values
            
        Returns:
            Dictionary containing regression metrics
        """
        from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
        
        metrics = {}
        
        # Basic metrics
        metrics["mse"] = mean_squared_error(y_true, y_pred)
        metrics["rmse"] = np.sqrt(metrics["mse"])
        metrics["mae"] = mean_absolute_error(y_true, y_pred)
        metrics["r2"] = r2_score(y_true, y_pred)
        
        # Additional metrics
        metrics["mean_error"] = np.mean(y_pred - y_true)
        metrics["std_error"] = np.std(y_pred - y_true)
        
        # Percentage errors
        mape = np.mean(np.abs((y_true - y_pred) / (y_true + 1e-8))) * 100
        metrics["mape"] = mape
        
        return metrics
    
    def calculate_embedding_metrics(self, embeddings: np.ndarray, 
                                  labels: np.ndarray) -> Dict[str, float]:
        """
        Calculate metrics specific to embedding models (e.g., speaker recognition).
        
        Args:
            embeddings: Embedding vectors
            labels: True labels
            
        Returns:
            Dictionary containing embedding metrics
        """
        from sklearn.metrics.pairwise import cosine_similarity
        
        metrics = {}
        
        # Calculate pairwise similarities
        similarities = cosine_similarity(embeddings)
        
        # Intra-class and inter-class similarities
        intra_class_sims = []
        inter_class_sims = []
        
        unique_labels = np.unique(labels)
        
        for i in range(len(embeddings)):
            for j in range(i + 1, len(embeddings)):
                sim = similarities[i, j]
                
                if labels[i] == labels[j]:
                    intra_class_sims.append(sim)
                else:
                    inter_class_sims.append(sim)
        
        metrics["intra_class_similarity_mean"] = np.mean(intra_class_sims)
        metrics["intra_class_similarity_std"] = np.std(intra_class_sims)
        metrics["inter_class_similarity_mean"] = np.mean(inter_class_sims)
        metrics["inter_class_similarity_std"] = np.std(inter_class_sims)
        
        # Separation metric
        separation = (np.mean(intra_class_sims) - np.mean(inter_class_sims))
        metrics["class_separation"] = separation
        
        # Equal Error Rate (EER) calculation
        # Create pairs and labels for verification task
        pairs = []
        pair_labels = []
        
        for i in range(len(embeddings)):
            for j in range(i + 1, len(embeddings)):
                pairs.append((i, j))
                pair_labels.append(1 if labels[i] == labels[j] else 0)
        
        # Calculate similarities for pairs
        pair_similarities = [similarities[i, j] for i, j in pairs]
        
        # Find EER threshold
        thresholds = np.linspace(min(pair_similarities), max(pair_similarities), 1000)
        best_eer = 1.0
        best_threshold = 0.0
        
        for threshold in thresholds:
            predictions = (np.array(pair_similarities) >= threshold).astype(int)
            
            # Calculate FAR and FRR
            true_pairs = np.array(pair_labels) == 1
            false_pairs = np.array(pair_labels) == 0
            
            if np.sum(true_pairs) > 0 and np.sum(false_pairs) > 0:
                far = np.sum(predictions[false_pairs]) / np.sum(false_pairs)  # False Accept Rate
                frr = 1 - np.sum(predictions[true_pairs]) / np.sum(true_pairs)  # False Reject Rate
                
                if abs(far - frr) < abs(best_eer - best_eer):
                    best_eer = (far + frr) / 2
                    best_threshold = threshold
        
        metrics["eer"] = best_eer
        metrics["eer_threshold"] = best_threshold
        
        return metrics
    
    def calculate_all_metrics(self, y_true: np.ndarray, y_pred: np.ndarray, 
                            y_prob: np.ndarray = None,
                            embeddings: np.ndarray = None) -> Dict[str, Any]:
        """
        Calculate all appropriate metrics based on task type.
        
        Args:
            y_true: True labels/values
            y_pred: Predicted labels/values
            y_prob: Prediction probabilities (for classification)
            embeddings: Embedding vectors (for embedding tasks)
            
        Returns:
            Dictionary containing all calculated metrics
        """
        if self.task_type == "classification":
            metrics = self.calculate_classification_metrics(y_true, y_pred, y_prob)
            
            # Add embedding metrics if available
            if embeddings is not None:
                embedding_metrics = self.calculate_embedding_metrics(embeddings, y_true)
                metrics.update(embedding_metrics)
            
            return metrics
        
        elif self.task_type == "regression":
            return self.calculate_regression_metrics(y_true, y_pred)
        
        else:
            raise ValueError(f"Unsupported task type: {self.task_type}")


class ModelProfiler:
    """
    Model performance profiler for analyzing computational requirements.
    
    Provides profiling of model inference time, memory usage, FLOPs,
    and other performance characteristics.
    """
    
    def __init__(self, model: nn.Module, device: str = "cpu"):
        """
        Initialize model profiler.
        
        Args:
            model: PyTorch model to profile
            device: Device for profiling
        """
        self.model = model.to(device)
        self.device = device
        self.model.eval()
    
    def profile_inference(self, input_tensor: torch.Tensor, 
                         num_runs: int = 100,
                         warmup_runs: int = 10) -> Dict[str, float]:
        """
        Profile model inference performance.
        
        Args:
            input_tensor: Sample input tensor
            num_runs: Number of inference runs for timing
            warmup_runs: Number of warmup runs
            
        Returns:
            Dictionary containing timing statistics
        """
        input_tensor = input_tensor.to(self.device)
        
        # Warmup runs
        with torch.no_grad():
            for _ in range(warmup_runs):
                _ = self.model(input_tensor)
        
        # Synchronize if using CUDA
        if self.device.startswith('cuda'):
            torch.cuda.synchronize()
        
        # Timing runs
        times = []
        with torch.no_grad():
            for _ in range(num_runs):
                start_time = time.perf_counter()
                output = self.model(input_tensor)
                
                if self.device.startswith('cuda'):
                    torch.cuda.synchronize()
                
                end_time = time.perf_counter()
                times.append(end_time - start_time)
        
        return {
            "mean_inference_time": np.mean(times),
            "std_inference_time": np.std(times),
            "min_inference_time": np.min(times),
            "max_inference_time": np.max(times),
            "median_inference_time": np.median(times),
            "throughput_samples_per_second": input_tensor.shape[0] / np.mean(times)
        }
    
    def profile_memory(self, input_tensor: torch.Tensor) -> Dict[str, float]:
        """
        Profile memory usage during inference.
        
        Args:
            input_tensor: Sample input tensor
            
        Returns:
            Dictionary containing memory statistics
        """
        input_tensor = input_tensor.to(self.device)
        
        if not self.device.startswith('cuda'):
            return {"error": "Memory profiling only available for CUDA devices"}
        
        # Clear cache and reset peak memory stats
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats(self.device)
        
        # Record initial memory
        initial_memory = torch.cuda.memory_allocated(self.device)
        
        # Run inference
        with torch.no_grad():
            output = self.model(input_tensor)
        
        # Record peak memory
        peak_memory = torch.cuda.max_memory_allocated(self.device)
        final_memory = torch.cuda.memory_allocated(self.device)
        
        return {
            "initial_memory_mb": initial_memory / (1024 * 1024),
            "peak_memory_mb": peak_memory / (1024 * 1024),
            "final_memory_mb": final_memory / (1024 * 1024),
            "memory_increase_mb": (peak_memory - initial_memory) / (1024 * 1024)
        }
    
    def count_parameters(self) -> Dict[str, int]:
        """Count model parameters."""
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        return {
            "total_parameters": total_params,
            "trainable_parameters": trainable_params,
            "non_trainable_parameters": total_params - trainable_params
        }
    
    def estimate_model_size(self) -> Dict[str, float]:
        """Estimate model size in different formats."""
        # Parameter size
        param_size = sum(p.numel() * p.element_size() for p in self.model.parameters())
        buffer_size = sum(b.numel() * b.element_size() for b in self.model.buffers())
        
        total_size = param_size + buffer_size
        
        return {
            "total_size_bytes": total_size,
            "total_size_mb": total_size / (1024 * 1024),
            "parameter_size_mb": param_size / (1024 * 1024),
            "buffer_size_mb": buffer_size / (1024 * 1024)
        }
    
    def generate_profile_report(self, input_tensor: torch.Tensor) -> str:
        """Generate comprehensive profiling report."""
        inference_stats = self.profile_inference(input_tensor)
        memory_stats = self.profile_memory(input_tensor)
        param_stats = self.count_parameters()
        size_stats = self.estimate_model_size()
        
        report = "# Model Profiling Report\n\n"
        
        # Model info
        report += "## Model Information\n"
        report += f"- Model: {self.model.__class__.__name__}\n"
        report += f"- Device: {self.device}\n"
        report += f"- Input shape: {input_tensor.shape}\n\n"
        
        # Parameters
        report += "## Parameters\n"
        report += f"- Total parameters: {param_stats['total_parameters']:,}\n"
        report += f"- Trainable parameters: {param_stats['trainable_parameters']:,}\n"
        report += f"- Non-trainable parameters: {param_stats['non_trainable_parameters']:,}\n\n"
        
        # Model size
        report += "## Model Size\n"
        report += f"- Total size: {size_stats['total_size_mb']:.2f} MB\n"
        report += f"- Parameter size: {size_stats['parameter_size_mb']:.2f} MB\n"
        report += f"- Buffer size: {size_stats['buffer_size_mb']:.2f} MB\n\n"
        
        # Inference performance
        report += "## Inference Performance\n"
        report += f"- Mean inference time: {inference_stats['mean_inference_time']*1000:.2f} ms\n"
        report += f"- Std inference time: {inference_stats['std_inference_time']*1000:.2f} ms\n"
        report += f"- Throughput: {inference_stats['throughput_samples_per_second']:.1f} samples/sec\n\n"
        
        # Memory usage
        if "error" not in memory_stats:
            report += "## Memory Usage\n"
            report += f"- Peak memory: {memory_stats['peak_memory_mb']:.2f} MB\n"
            report += f"- Memory increase: {memory_stats['memory_increase_mb']:.2f} MB\n\n"
        
        return report