dataset.py

import numpy as np
import torch
import torch.utils.data as data
import matplotlib.pyplot as plt
from typing import Union, Tuple, Optional


class SmilingFaceDataset:
    """
    A 2D probability distribution in the shape of a smiling face.
    Features two circular eye modes and a half-moon mouth mode.
    Optimized for PyTorch diffusion model training.
    """
    
    def __init__(self, device: str = 'cpu', dtype: torch.dtype = torch.float32, seed: Optional[int] = 42):
        """
        Initialize the smiling face distribution.
        
        Args:
            device: PyTorch device ('cpu' or 'cuda')
            dtype: PyTorch data type
            seed: Random seed for reproducible sampling
        """
        self.device = device
        self.dtype = dtype
        self.seed = seed
        
        # Set random seeds for reproducibility
        if seed is not None:
            np.random.seed(seed)
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(seed)
                torch.cuda.manual_seed_all(seed)
        
        # Parameters for eyes
        self.eye_distance = 0.2  # distance from center
        self.eye_height = 0.2    # height above center
        self.eye_sigma = 0.04    # eye width
        self.eye_ratio = 0.4     # height ratio between eyes and mouth (after normalization)
        
        # Parameters for mouth arc
        self.mouth_center = (0.0, 0.0)
        self.mouth_radius = 0.3
        self.mouth_theta1 = 7 * np.pi / 6  # smile arc start
        self.mouth_theta2 = 11 * np.pi / 6  # smile arc end
        self.mouth_sigma = self.eye_sigma  # mouth thickness
        self.mouth_n_points = 100  # discretization points
        
        # Domain bounds
        self.x_min, self.x_max = -0.5, 0.5
        self.y_min, self.y_max = -0.5, 0.5
        
        # Precompute mouth arc points
        self._precompute_mouth_arc()
        
        # Compute normalization factors
        self._compute_normalization()
    
    def _precompute_mouth_arc(self):
        """Precompute the mouth arc points for efficient sampling."""
        thetas = np.linspace(self.mouth_theta1, self.mouth_theta2, self.mouth_n_points)
        self.arc_x = self.mouth_center[0] + self.mouth_radius * np.cos(thetas)
        self.arc_y = self.mouth_center[1] + self.mouth_radius * np.sin(thetas)
        self.dtheta = (self.mouth_theta2 - self.mouth_theta1) / (self.mouth_n_points - 1)
    
    def _eye_component(self, x: np.ndarray, y: np.ndarray, x0: float, y0: float) -> np.ndarray:
        """Compute eye component (Gaussian bump)."""
        return np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * self.eye_sigma**2))
    
    def _mouth_component(self, x: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Compute mouth component (sum of Gaussians along arc)."""
        val = np.zeros_like(x)
        for px, py in zip(self.arc_x, self.arc_y):
            val += np.exp(-((x - px)**2 + (y - py)**2) / (2 * self.mouth_sigma**2))
        return val * self.dtheta
    
    def _compute_normalization(self):
        """Compute normalization factors to make all components have equal peak height."""
        # Create evaluation grid
        n_grid = 200
        xs = np.linspace(self.x_min, self.x_max, n_grid)
        ys = np.linspace(self.y_min, self.y_max, n_grid)
        X, Y = np.meshgrid(xs, ys)
        
        # Compute individual components
        left_eye = self._eye_component(X, Y, -self.eye_distance, self.eye_height)
        right_eye = self._eye_component(X, Y, self.eye_distance, self.eye_height)
        mouth = self._mouth_component(X, Y)
        
        # Find maximum values
        left_eye_max = np.max(left_eye)
        right_eye_max = np.max(right_eye)
        mouth_max = np.max(mouth)
        
        # Normalize to global maximum
        global_max = max(left_eye_max, right_eye_max, mouth_max)
        self.left_eye_factor = global_max / left_eye_max
        self.right_eye_factor = global_max / right_eye_max * self.eye_ratio
        self.mouth_factor = global_max / mouth_max
        
        # Compute overall normalization constant
        unnormalized_pdf = (left_eye * self.left_eye_factor + 
                          right_eye * self.right_eye_factor + 
                          mouth * self.mouth_factor)
        self.Z = np.trapezoid(np.trapezoid(unnormalized_pdf, xs), ys)
    
    def pdf(self, x: Union[np.ndarray, torch.Tensor], y: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
        """
        Evaluate the probability density function.
        
        Args:
            x, y: Coordinates (numpy arrays or torch tensors)
            
        Returns:
            PDF values at given coordinates
        """
        is_tensor = isinstance(x, torch.Tensor)
        
        if is_tensor:
            x_np = x.cpu().numpy()
            y_np = y.cpu().numpy()
        else:
            x_np, y_np = x, y
        
        # Compute normalized components
        left_eye = self._eye_component(x_np, y_np, -self.eye_distance, self.eye_height) * self.left_eye_factor
        right_eye = self._eye_component(x_np, y_np, self.eye_distance, self.eye_height) * self.right_eye_factor
        mouth = self._mouth_component(x_np, y_np) * self.mouth_factor
        
        # Combined normalized PDF
        pdf_vals = (left_eye + right_eye + mouth) / self.Z
        
        if is_tensor:
            return torch.tensor(pdf_vals, dtype=self.dtype, device=self.device)
        return pdf_vals
    
    def sample(self, n_samples: int, return_tensor: bool = True, seed: Optional[int] = None) -> Union[torch.Tensor, np.ndarray]:
        """
        Sample points from the smiling face distribution using rejection sampling.
        
        Args:
            n_samples: Number of samples to generate
            return_tensor: If True, return torch.Tensor; if False, return numpy array
            seed: Optional seed for this specific sampling (overrides instance seed)
            
        Returns:
            Samples with shape (n_samples, 2)
        """
        # Set seed if provided
        if seed is not None:
            np.random.seed(seed)
        
        # Estimate maximum PDF value for rejection sampling
        max_pdf_estimate = 3.0 / self.Z  # Conservative upper bound
        
        samples = []
        n_generated = 0
        
        while n_generated < n_samples:
            # Generate candidate points
            batch_size = min(n_samples * 3, 10000)  # Generate in batches with some overhead
            x_candidates = np.random.uniform(self.x_min, self.x_max, batch_size)
            y_candidates = np.random.uniform(self.y_min, self.y_max, batch_size)
            
            # Evaluate PDF
            pdf_vals = self.pdf(x_candidates, y_candidates)
            
            # Acceptance probabilities
            accept_probs = pdf_vals / max_pdf_estimate
            accept_mask = np.random.uniform(0, 1, batch_size) < accept_probs
            
            # Accept samples
            accepted_x = x_candidates[accept_mask]
            accepted_y = y_candidates[accept_mask]
            
            if len(accepted_x) > 0:
                batch_samples = np.column_stack([accepted_x, accepted_y])
                samples.append(batch_samples)
                n_generated += len(batch_samples)
        
        # Combine and truncate to exact number requested
        all_samples = np.vstack(samples)[:n_samples]
        
        if return_tensor:
            return torch.tensor(all_samples, dtype=self.dtype, device=self.device)
        return all_samples
    
    def to_tensor(self, data: np.ndarray) -> torch.Tensor:
        """Convert numpy array to torch tensor with correct dtype and device."""
        return torch.tensor(data, dtype=self.dtype, device=self.device)
    
    def get_bounds(self) -> Tuple[float, float, float, float]:
        """Get the domain bounds (x_min, x_max, y_min, y_max)."""
        return self.x_min, self.x_max, self.y_min, self.y_max
    
    def plot_pdf_contour(self, n_grid: int = 200, figsize: Tuple[int, int] = (8, 6), 
                        levels: int = 50, title: str = "Smiling Face PDF") -> None:
        """
        Plot the PDF as a contour plot.
        
        Args:
            n_grid: Number of grid points per dimension
            figsize: Figure size
            levels: Number of contour levels
            title: Plot title
        """
        # Create evaluation grid
        xs = np.linspace(self.x_min, self.x_max, n_grid)
        ys = np.linspace(self.y_min, self.y_max, n_grid)
        X, Y = np.meshgrid(xs, ys)
        
        # Evaluate PDF
        PDF_vals = self.pdf(X, Y)
        
        # Create plot
        plt.figure(figsize=figsize)
        contour = plt.contourf(X, Y, PDF_vals, levels=levels, cmap='viridis')
        plt.colorbar(contour, label='PDF Value')
        plt.xlabel('X coordinate')
        plt.ylabel('Y coordinate')
        plt.title(title)
        plt.xlim(self.x_min, self.x_max)
        plt.ylim(self.y_min, self.y_max)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig('pdf_contour.png')
    
    def plot_samples(self, samples: Union[torch.Tensor, np.ndarray], 
                    figsize: Tuple[int, int] = (8, 6), alpha: float = 0.6, 
                    title: str = "Sampled Points", s: int = 20) -> None:
        """
        Plot sampled points as a scatter plot.
        
        Args:
            samples: Sample points with shape (n_samples, 2)
            figsize: Figure size
            alpha: Point transparency
            title: Plot title
            s: Point size
        """
        # Convert to numpy if needed
        if isinstance(samples, torch.Tensor):
            samples_np = samples.cpu().numpy()
        else:
            samples_np = samples
        
        # Create plot
        plt.figure(figsize=figsize)
        plt.scatter(samples_np[:, 0], samples_np[:, 1], 
                   alpha=alpha, s=s, c='blue', edgecolors='black', linewidth=0.5)
        plt.xlabel('X coordinate')
        plt.ylabel('Y coordinate')
        plt.title(f'{title} ({len(samples_np)} points)')
        plt.xlim(self.x_min, self.x_max)
        plt.ylim(self.y_min, self.y_max)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig('samples.png')
    
    def plot_comparison(self, samples: Union[torch.Tensor, np.ndarray], 
                       n_grid: int = 200, figsize: Tuple[int, int] = (16, 6)) -> None:
        """
        Plot PDF contour and samples side by side for comparison.
        
        Args:
            samples: Sample points with shape (n_samples, 2)
            n_grid: Number of grid points for PDF evaluation
            figsize: Figure size
        """
        # Convert to numpy if needed
        if isinstance(samples, torch.Tensor):
            samples_np = samples.cpu().numpy()
        else:
            samples_np = samples
        
        # Create evaluation grid for PDF
        xs = np.linspace(self.x_min, self.x_max, n_grid)
        ys = np.linspace(self.y_min, self.y_max, n_grid)
        X, Y = np.meshgrid(xs, ys)
        PDF_vals = self.pdf(X, Y)
        
        # Create subplots
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
        
        # Plot 1: PDF contour
        contour = ax1.contourf(X, Y, PDF_vals, levels=50, cmap='viridis')
        fig.colorbar(contour, ax=ax1, label='PDF Value')
        ax1.set_xlabel('X coordinate')
        ax1.set_ylabel('Y coordinate')
        ax1.set_title('True PDF')
        ax1.set_xlim(self.x_min, self.x_max)
        ax1.set_ylim(self.y_min, self.y_max)
        ax1.set_aspect('equal', adjustable='box')
        ax1.grid(True, alpha=0.3)

        # Plot 2: Samples
        ax2.scatter(samples_np[:, 0], samples_np[:, 1], 
                   alpha=0.6, s=15, c='blue', edgecolors='black', linewidth=0.3)
        ax2.set_xlabel('X coordinate')
        ax2.set_ylabel('Y coordinate')
        ax2.set_title(f'Samples ({len(samples_np)} points)')
        ax2.set_xlim(self.x_min, self.x_max)
        ax2.set_ylim(self.y_min, self.y_max)
        ax2.set_aspect('equal', adjustable='box')
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('comparison.png')

    def get_dataloader(self, n_samples: int, batch_size: int = 64, shuffle: bool = True, 
                      num_workers: int = 0, pin_memory: bool = True, seed: Optional[int] = None) -> data.DataLoader:
        """
        Create a PyTorch DataLoader for training.
        
        Args:
            n_samples: Number of samples to generate for the dataset
            batch_size: Batch size for training
            shuffle: Whether to shuffle the data
            num_workers: Number of worker processes for data loading
            pin_memory: Whether to pin memory for faster GPU transfer
            seed: Optional seed for sample generation (uses instance seed if None)
            
        Returns:
            PyTorch DataLoader
        """
        # Use provided seed or instance seed
        sample_seed = seed if seed is not None else self.seed
        
        # Generate samples
        samples = self.sample(n_samples, return_tensor=True, seed=sample_seed)
        
        # Create TensorDataset
        dataset = data.TensorDataset(samples)
        
        # Create DataLoader
        dataloader = data.DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=num_workers,
            pin_memory=pin_memory and torch.cuda.is_available()
        )
        
        return dataloader
    
    def reset_seed(self, seed: Optional[int] = None):
        """
        Reset the random seed for reproducible sampling.
        
        Args:
            seed: New seed value (uses original seed if None)
        """
        if seed is not None:
            self.seed = seed
        np.random.seed(self.seed)
        
    def true_score(self, x: Union[np.ndarray, torch.Tensor], y: Union[np.ndarray, torch.Tensor],
                   eps: float = 1e-4, pdf_threshold: float = 1e-6) -> Union[np.ndarray, torch.Tensor]:
        """
        Compute the true score function ∇_x log p(x) = ∇_x (log p(x)).
        
        The score function is the gradient of the log probability density function.
        This is computed using numerical differentiation, but only at points where
        the PDF is above a threshold to avoid numerical instability.
        
        Args:
            x: X coordinates
            y: Y coordinates  
            eps: Small value for numerical differentiation
            pdf_threshold: Only compute score where PDF > threshold
            
        Returns:
            Score function values with shape (..., 2) where the last dimension
            contains [∂log p/∂x, ∂log p/∂y]. Score is zero where PDF < threshold.
        """
        # Convert to numpy if needed
        is_tensor = isinstance(x, torch.Tensor)
        if is_tensor:
            x_np = x.detach().cpu().numpy()
            y_np = y.detach().cpu().numpy()
        else:
            x_np = x
            y_np = y
        
        # For regular grids, use np.gradient which is more robust
        if x_np.ndim == 2 and y_np.ndim == 2:  # Grid input
            # Compute PDF on the grid
            pdf_values = self.pdf(x_np, y_np)
            
            # Add small epsilon to avoid log(0)
            log_pdf = np.log(pdf_values + 1e-8)
            
            # Use np.gradient to compute the score (like in your example)
            # np.gradient returns [grad_y, grad_x] for 2D arrays
            grad_y, grad_x = np.gradient(log_pdf, 
                                        y_np[1, 0] - y_np[0, 0],  # dy spacing
                                        x_np[0, 1] - x_np[0, 0])  # dx spacing
            
            score_x = grad_x
            score_y = grad_y
            
        else:  # Point-wise input - use finite differences
            # Compute PDF at current points to check threshold
            pdf_current = self.pdf(x_np, y_np)
            
            # Only compute scores where PDF is above threshold
            valid_mask = pdf_current > pdf_threshold
            
            # Initialize score arrays
            score_x = np.zeros_like(pdf_current)
            score_y = np.zeros_like(pdf_current)
            
            if np.any(valid_mask):
                # Compute log PDF only where needed
                log_pdf_x_plus = np.log(self.pdf(x_np + eps, y_np) + 1e-8)
                log_pdf_x_minus = np.log(self.pdf(x_np - eps, y_np) + 1e-8)
                log_pdf_y_plus = np.log(self.pdf(x_np, y_np + eps) + 1e-8)
                log_pdf_y_minus = np.log(self.pdf(x_np, y_np - eps) + 1e-8)
                
                # Score function: ∇_x log p(x) computed directly
                score_x_all = (log_pdf_x_plus - log_pdf_x_minus) / (2 * eps)
                score_y_all = (log_pdf_y_plus - log_pdf_y_minus) / (2 * eps)
                
                # Only keep scores where PDF is above threshold
                score_x[valid_mask] = score_x_all[valid_mask]
                score_y[valid_mask] = score_y_all[valid_mask]
        
        # Stack into score vectors
        if x_np.ndim == 0:  # scalar inputs
            score = np.array([score_x, score_y])
        else:  # array inputs
            score = np.stack([score_x, score_y], axis=-1)
        
        # Convert back to tensor if needed
        if is_tensor:
            return torch.tensor(score, dtype=self.dtype, device=self.device)
        return score
    
    def plot_score_field(self, n_grid: int = 30, figsize: Tuple[int, int] = (8, 6),
                        scale: float = 20, title: str = "True Score Field") -> None:
        """
        Plot the true score field as a vector field.
        
        Args:
            n_grid: Number of grid points per dimension
            figsize: Figure size
            scale: Scale factor for arrow size
            title: Plot title
        """
        # Create evaluation grid
        xs = np.linspace(self.x_min, self.x_max, n_grid)
        ys = np.linspace(self.y_min, self.y_max, n_grid)
        X, Y = np.meshgrid(xs, ys)
        
        # Compute score field
        scores = self.true_score(X, Y)
        
        # Create plot
        plt.figure(figsize=figsize)
        plt.quiver(X, Y, scores[:, :, 0], scores[:, :, 1], 
                  alpha=0.7, scale=scale, width=0.003)
        plt.xlabel('X coordinate')
        plt.ylabel('Y coordinate')
        plt.title(title)
        plt.xlim(self.x_min, self.x_max)
        plt.ylim(self.y_min, self.y_max)
        plt.gca().set_aspect('equal', adjustable='box')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig('score_field.png')

    def _create_score_function(self):
        """
        Get the true score function as a reusable callable.
        
        This returns a function that can be used for both training (loss computation)
        and sampling (Langevin dynamics).
        
        Returns:
            Callable that takes x (tensor) and returns scores (tensor)
        """
        def score_function(x: torch.Tensor) -> torch.Tensor:
            """
            Compute true scores for input points.
            
            Args:
                x: Input points (batch_size, 2)
                
            Returns:
                True scores (batch_size, 2)
            """
            # Convert to numpy for computation
            x_np = x.detach().cpu().numpy()
            
            # Compute true scores
            scores_np = self.true_score(x_np[:, 0], x_np[:, 1])
            
            # Convert back to tensor with same device/dtype as input
            if isinstance(scores_np, torch.Tensor):
                return scores_np.to(x.device, dtype=x.dtype)
            else:
                return torch.tensor(scores_np, device=x.device, dtype=x.dtype)
        
        return score_function
    
    def get_true_score_function(self):
        """Get a function that computes true scores for given input points."""
        return self._create_score_function()
    
    def true_score_with_noise(self, x: Union[np.ndarray, torch.Tensor], y: Union[np.ndarray, torch.Tensor],
                             sigma: float, eps: float = 1e-4) -> Union[np.ndarray, torch.Tensor]:
        """
        Compute the score function when Gaussian noise with variance σ² is added to the distribution.
        
        For a mixture of Gaussians, adding noise σ² means each component's variance increases by σ².
        This function recomputes the score field for the noise-modified distribution.
        
        Args:
            x: X coordinates
            y: Y coordinates  
            sigma: Noise level (standard deviation) to add to each component
            eps: Small value for numerical differentiation
            
        Returns:
            Score function values for the noisy distribution with shape (..., 2)
        """
        # Convert to numpy if needed
        is_tensor = isinstance(x, torch.Tensor)
        if is_tensor:
            x_np = x.detach().cpu().numpy()
            y_np = y.detach().cpu().numpy()
        else:
            x_np, y_np = x, y
        
        # Compute score using finite differences with the noisy PDF
        # PDF at neighboring points
        pdf_x_plus = self.pdf_with_noise(x_np + eps, y_np, sigma)
        pdf_x_minus = self.pdf_with_noise(x_np - eps, y_np, sigma)
        pdf_y_plus = self.pdf_with_noise(x_np, y_np + eps, sigma)
        pdf_y_minus = self.pdf_with_noise(x_np, y_np - eps, sigma)
        
        # Compute score as gradient of log PDF
        score_x = (np.log(pdf_x_plus + 1e-8) - np.log(pdf_x_minus + 1e-8)) / (2 * eps)
        score_y = (np.log(pdf_y_plus + 1e-8) - np.log(pdf_y_minus + 1e-8)) / (2 * eps)
        
        # Stack into score vectors
        if x_np.ndim == 0:  # scalar inputs
            score = np.array([score_x, score_y])
        else:  # array inputs
            score = np.stack([score_x, score_y], axis=-1)
        
        # Convert back to tensor if needed
        if is_tensor:
            return torch.tensor(score, dtype=self.dtype, device=self.device)
        return score
    
    def pdf_with_noise(self, x: Union[np.ndarray, torch.Tensor], y: Union[np.ndarray, torch.Tensor],
                      sigma: float) -> Union[np.ndarray, torch.Tensor]:
        """
        Evaluate the PDF when Gaussian noise with variance σ² is added to the distribution.
        
        For a mixture of Gaussians, adding noise σ² means each component's variance increases by σ².
        This function recomputes the PDF for the noise-modified distribution.
        
        Args:
            x, y: Coordinates (numpy arrays or torch tensors)
            sigma: Noise level (standard deviation) to add to each component
            
        Returns:
            PDF values for the noisy distribution at given coordinates
        """
        is_tensor = isinstance(x, torch.Tensor)
        
        if is_tensor:
            x_np = x.cpu().numpy()
            y_np = y.cpu().numpy()
        else:
            x_np, y_np = x, y
        
        # New variances with added noise
        new_eye_var = self.eye_sigma**2 + sigma**2
        new_mouth_var = self.mouth_sigma**2 + sigma**2
        
        new_eye_sigma = np.sqrt(new_eye_var)
        new_mouth_sigma = np.sqrt(new_mouth_var)
        
        # Helper functions for noisy components
        def _noisy_eye_component(x, y, x0, y0, sigma_noisy):
            return np.exp(-((x - x0)**2 + (y - y0)**2) / (2 * sigma_noisy**2))
        
        def _noisy_mouth_component(x, y, sigma_noisy):
            val = np.zeros_like(x)
            for px, py in zip(self.arc_x, self.arc_y):
                val += np.exp(-((x - px)**2 + (y - py)**2) / (2 * sigma_noisy**2))
            return val * self.dtheta
        
        # Compute normalized components with noisy variances
        left_eye = _noisy_eye_component(x_np, y_np, -self.eye_distance, self.eye_height, new_eye_sigma) * self.left_eye_factor
        right_eye = _noisy_eye_component(x_np, y_np, self.eye_distance, self.eye_height, new_eye_sigma) * self.right_eye_factor
        mouth = _noisy_mouth_component(x_np, y_np, new_mouth_sigma) * self.mouth_factor
        
        # Combined normalized PDF (using original normalization constant for approximation)
        pdf_vals = (left_eye + right_eye + mouth) / self.Z
        
        if is_tensor:
            return torch.tensor(pdf_vals, dtype=self.dtype, device=self.device)
        return pdf_vals


if __name__ == "__main__":
    """
    Create and visualize the smiling face dataset when run directly.
    """
    print("🎨 Smiling Face Dataset")
    print("=" * 40)
    
    # Create dataset
    dataset = SmilingFaceDataset(device='cpu', seed=42)
    
    print(f"Dataset bounds: {dataset.get_bounds()}")
    print("Creating overview plot...")
    
    # Import and use the visualization function
    try:
        from vis import plot_dataset_overview
        
        # Create the overview plot
        save_path = plot_dataset_overview(
            dataset=dataset,
            n_samples=2000,
            n_grid=200,
            show=False  # Don't show in headless environments
        )
        
        print(f"✅ Dataset overview saved to: {save_path}")
        
    except ImportError:
        print("⚠️  vis module not available, creating basic plots...")
        
        # Fallback to basic plotting
        samples = dataset.sample(2000, return_tensor=False)
        
        # Create basic plots
        dataset.plot_pdf_contour(title="Smiling Face PDF Distribution")
        dataset.plot_samples(samples, title="Sampled Points from Distribution")
        dataset.plot_comparison(samples)
        
        print("✅ Basic plots created")
    
    print("🚀 Dataset ready for PyTorch training!")