import torch

class DebugVisualizer:
    """
    Debug visualization tools for light estimation analysis.
    """
    
    @staticmethod
    def generate_debug_mask(mask):
        """
        Generate debug visualization of the binary mask.
        """
        # Ensure mask is 3D (B, H, W)
        if mask.dim() == 4 and mask.shape[-1] == 1:
            mask = mask.squeeze(-1)

        # Convert to RGB visualization: white for lit areas, black for unlit
        debug_mask = mask.float().unsqueeze(-1).repeat(1, 1, 1, 3)
        return debug_mask

    @staticmethod
    def generate_lit_normals_visualization(normals, mask):
        """
        Generate visualization showing only the lit normal vectors.
        """
        normal_rgb = (normals + 1.0) / 2.0
        lit_mask = mask.unsqueeze(-1)
        lit_normals_viz = normal_rgb * lit_mask
        return lit_normals_viz

    @staticmethod
    def generate_directional_visualization(normals, results):
        """
        Generate color-coded visualization of directional analysis.
        """
        batch_size, height, width, _ = normals.shape
        viz = torch.zeros(batch_size, height, width, 3, device=normals.device)

        x_category = results['x_category']
        y_category = results['y_category']
        hard_soft_index = results['hard_soft_index']
        overall_confidence = results['confidence']['overall_confidence']

        # Base colors
        if x_category == "left":
            base_color = torch.tensor([255, 100, 100], device=normals.device)
        elif x_category == "right":
            base_color = torch.tensor([200, 80, 80], device=normals.device)
        else:
            base_color = torch.tensor([150, 150, 150], device=normals.device)

        # Adjust for Y direction
        if y_category == "top":
            base_color += torch.tensor([0, 100, 50], device=normals.device)
        elif y_category == "bottom":
            base_color += torch.tensor([50, -50, 0], device=normals.device)

        # Apply softness and confidence
        softness_factor = 1.0 - hard_soft_index
        final_color = base_color * (0.7 + 0.3 * softness_factor) * overall_confidence
        final_color = torch.clamp(final_color, 0, 255)

        viz[:, :, :] = final_color.byte()
        return viz.float() / 255.0
    
    @staticmethod
    def generate_colormap_preview(height=256, width=256):
        """
        Generate a visual preview of the colormap used for analysis.
        """
        x = torch.linspace(-1, 1, width)
        y = torch.linspace(-1, 1, height)
        xx, yy = torch.meshgrid(y, x, indexing='ij')
        
        # Mock results for visualization
        mock_normals = torch.stack([xx, yy, torch.zeros_like(xx)], dim=-1).unsqueeze(0)
        
        # Generate a color for each point in the grid
        colormap = torch.zeros(height, width, 3)
        for i in range(height):
            for j in range(width):
                # Pseudo-analysis for color mapping
                x_val, y_val = xx[i, j], yy[i, j]
                
                # Determine color based on position
                red = 128 + 127 * x_val
                green = 128 + 127 * y_val
                blue = 128 - 127 * abs(x_val)
                
                colormap[i, j] = torch.tensor([red, green, blue])
        
        return colormap.unsqueeze(0).byte().float() / 255.0

    @staticmethod
    def generate_color_histogram(normals, title="Color Distribution"):
        """
        Generate a color histogram visualization of normal vectors.

        Args:
            normals: Normal vectors tensor (B, H, W, 3)
            title: Title for the histogram

        Returns:
            histogram_viz: RGB visualization of color distribution (1, 256, 256, 3)
        """
        # Flatten normals to get color distribution
        flat_normals = normals.view(-1, 3) * 255 + 128  # Scale to 0-255 range
        flat_normals = torch.clamp(flat_normals, 0, 255)

        # Create histogram visualization
        hist_size = 256
        histogram = torch.zeros(hist_size, hist_size, 3)

        # Create a simple 2D histogram representation
        for i in range(min(flat_normals.shape[0], 10000)):  # Sample for performance
            x = int(flat_normals[i, 0].item())
            y = int(flat_normals[i, 1].item())
            if 0 <= x < hist_size and 0 <= y < hist_size:
                # Color based on Z component (blue channel)
                intensity = flat_normals[i, 2].item() / 255.0
                histogram[x, y] = torch.tensor([intensity, intensity, 1.0 - intensity])

        return histogram.unsqueeze(0).byte().float() / 255.0