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 clear reference guide for normal map color encoding.

        This shows how 3D surface orientations are encoded as RGB colors
        in normal maps, making it easy to understand what each color means.
        """
        # Create reference visualization
        viz = torch.zeros(height, width, 3)

        # Color legend for normal map encoding
        normal_colors = {
            'right': torch.tensor([1.0, 0.5, 0.5]),      # Red - Right facing
            'left': torch.tensor([0.5, 1.0, 1.0]),       # Cyan - Left facing
            'up': torch.tensor([0.5, 1.0, 0.5]),         # Green - Up facing
            'down': torch.tensor([1.0, 0.5, 1.0]),       # Purple - Down facing
            'out': torch.tensor([0.5, 0.5, 1.0]),        # Blue - Outward facing
            'flat': torch.tensor([0.5, 0.5, 0.5]),       # Gray - Flat/neutral
        }

        # Create visual layout
        section_height = height // 3
        section_width = width // 3

        sections = [
            ('X+ (Right)', normal_colors['right'], 0, 0),
            ('X- (Left)', normal_colors['left'], 0, 1),
            ('Y+ (Up)', normal_colors['up'], 1, 0),
            ('Y- (Down)', normal_colors['down'], 1, 1),
            ('Z+ (Out)', normal_colors['out'], 2, 0),
            ('Z=0 (Flat)', normal_colors['flat'], 2, 1),
        ]

        # Draw color sections with labels
        for label, color, row, col in sections:
            y_start, x_start = row * section_height, col * section_width

            # Fill section with color
            viz[y_start:y_start+section_height, x_start:x_start+section_width] = color

            # Add simple label area (could be enhanced with text rendering)
            label_y = y_start + section_height // 2
            label_x = x_start + 10

        # Add title and explanation
        title_area_height = height // 6
        viz[:title_area_height, :] = 0.8  # Light gray background for title

        return viz.unsqueeze(0)

    @staticmethod
    def generate_color_histogram(normals, title="Normal Distribution"):
        """
        Generate a clear, informative histogram of normal map color distribution.

        Args:
            normals: Normal vectors tensor (B, H, W, 3) in [-1, 1] range
            title: Title for the histogram

        Returns:
            histogram_viz: Clear visualization with labels and reference (1, 512, 512, 3)
        """
        return DebugVisualizer._create_enhanced_histogram(normals, title, show_reference=True)

    @staticmethod
    def generate_histogram_before(normals, title="Before Masking"):
        """
        Generate histogram of full normal map distribution.

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

        Returns:
            histogram_viz: RGB histogram of full distribution (1, 256, 256, 3)
        """
        return DebugVisualizer._create_color_histogram(normals, title)

    @staticmethod
    def generate_histogram_after(lit_normals, title="After Masking (Lit Areas)"):
        """
        Generate histogram of lit areas normal distribution.

        Args:
            lit_normals: Normal vectors tensor (N, 3) for lit areas only
            title: Title for the histogram

        Returns:
            histogram_viz: RGB histogram of lit areas distribution (1, 256, 256, 3)
        """
        # Add batch dimension if needed for consistent processing
        if lit_normals.dim() == 2:
            lit_normals = lit_normals.unsqueeze(0)

        return DebugVisualizer._create_color_histogram(lit_normals, title)

    @staticmethod
    def _create_color_histogram(normals, title="Color Distribution"):
        """
        Internal method to create color histogram from normal vectors.

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

        Returns:
            histogram_viz: RGB visualization (1, 256, 256, 3)
        """
        # Flatten normals to analyze distribution
        flat_normals = normals.view(-1, 3)

        if flat_normals.shape[0] == 0:
            return torch.zeros(1, 256, 256, 3)

        # Scale normal vectors from [-1, 1] to [0, 255] for histogram bins
        scaled_normals = (flat_normals + 1.0) * 127.5
        scaled_normals = torch.clamp(scaled_normals, 0, 255)

        # Create 2D histogram using Red and Green channels as coordinates
        hist_size = 256
        histogram = torch.zeros(hist_size, hist_size)

        # Sample for performance
        sample_size = min(flat_normals.shape[0], 100000)
        indices = torch.randperm(flat_normals.shape[0])[:sample_size]

        for i in indices:
            normal = scaled_normals[i]
            r, g, b = normal[0].item(), normal[1].item(), normal[2].item()

            # Use red and green as histogram coordinates
            r_bin = int(r)
            g_bin = int(g)

            if 0 <= r_bin < hist_size and 0 <= g_bin < hist_size:
                histogram[r_bin, g_bin] += 1.0

        # Normalize for visualization
        if histogram.max() > 0:
            histogram = histogram / histogram.max()

        # Create RGB visualization
        rgb_histogram = torch.zeros(hist_size, hist_size, 3)

        for i in range(hist_size):
            for j in range(hist_size):
                density = histogram[i, j].item()

                if density > 0:
                    # Color by position and density
                    # X direction (red) + Y direction (green) + density (blue)
                    x_intensity = i / 255.0  # Red channel shows X position
                    y_intensity = j / 255.0  # Green channel shows Y position
                    z_intensity = density    # Blue channel shows density

                    rgb_histogram[i, j] = torch.tensor([
                        x_intensity,  # X direction indicator
                        y_intensity,  # Y direction indicator
                        z_intensity   # Density indicator
                    ])

        return rgb_histogram.unsqueeze(0)

    @staticmethod
    def _create_enhanced_histogram(normals, title="Normal Distribution", show_reference=True):
        """
        Create an enhanced histogram with labels, legends, and clear visual design.

        Args:
            normals: Normal vectors tensor (B, H, W, 3) or (N, 3)
            title: Title for the histogram
            show_reference: Whether to show reference information

        Returns:
            enhanced_viz: Complete visualization with labels (1, 512, 512, 3)
        """
        # Get base histogram (256x256)
        base_histogram = DebugVisualizer._create_color_histogram_internal(normals)

        if base_histogram is None:
            return torch.zeros(1, 512, 512, 3)

        # Create enhanced visualization with labels and reference
        viz_size = 512
        hist_size = 256
        enhanced = torch.zeros(viz_size, viz_size, 3)

        # Place histogram in top-left area (256x256)
        enhanced[:hist_size, :hist_size] = base_histogram[0]

        # Add reference information (right side)
        if show_reference:
            enhanced = DebugVisualizer._add_histogram_labels(enhanced, title)

        return enhanced.unsqueeze(0)

    @staticmethod
    def _create_color_histogram_internal(normals):
        """Internal histogram creation without labels"""
        flat_normals = normals.view(-1, 3)
        if flat_normals.shape[0] == 0:
            return None

        # Scale to [0, 255] for histogram bins
        scaled_normals = (flat_normals + 1.0) * 127.5
        scaled_normals = torch.clamp(scaled_normals, 0, 255)

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

        sample_size = min(flat_normals.shape[0], 50000)
        indices = torch.randperm(flat_normals.shape[0])[:sample_size]

        for i in indices:
            normal = scaled_normals[i]
            r, g, b = normal[0].item(), normal[1].item(), normal[2].item()
            r_bin, g_bin = int(r), int(g)

            if 0 <= r_bin < hist_size and 0 <= g_bin < hist_size:
                histogram[r_bin, g_bin] += 1.0

        # Normalize
        if histogram.max() > 0:
            histogram = histogram / histogram.max()

        # Create visualization with clear color mapping
        rgb_histogram = torch.zeros(hist_size, hist_size, 3)

        # Add coordinate grid lines
        grid_color = torch.tensor([0.3, 0.3, 0.3])
        for i in range(0, hist_size, 32):
            rgb_histogram[i, :, :] = grid_color
            rgb_histogram[:, i, :] = grid_color

        # Add density visualization
        for i in range(hist_size):
            for j in range(hist_size):
                density = histogram[i, j].item()
                if density > 0.01:  # Threshold to reduce noise
                    # Heatmap: blue (low) to red (high)
                    if density < 0.33:
                        intensity = density * 3
                        rgb_histogram[i, j] = torch.tensor([0.0, 0.0, intensity])
                    elif density < 0.66:
                        intensity = (density - 0.33) * 3
                        rgb_histogram[i, j] = torch.tensor([intensity, 0.0, 0.0])
                    else:
                        intensity = (density - 0.66) * 3
                        rgb_histogram[i, j] = torch.tensor([1.0, intensity, 0.0])

        return rgb_histogram.unsqueeze(0)

    @staticmethod
    def _add_histogram_labels(enhanced_viz, title):
        """Add labels and reference information to histogram"""
        viz_size = enhanced_viz.shape[0]

        # Add title (top area)
        title_y = 20
        # Simple title for now - could be enhanced with text rendering

        # Add axis labels
        # X-axis: "Left ← X Direction → Right"
        # Y-axis: "Bottom ← Y Direction → Top"

        # Add color legend (right side)
        legend_x = int(viz_size * 0.8)
        legend_width = viz_size - legend_x

        # Color legend: show what colors mean
        legend_colors = [
            torch.tensor([0.0, 0.0, 1.0]),  # Blue = low density
            torch.tensor([0.0, 1.0, 0.0]),  # Green = medium density
            torch.tensor([1.0, 0.0, 0.0]),  # Red = high density
        ]

        for i, color in enumerate(legend_colors):
            y = int(viz_size * 0.3) + i * 30
            enhanced_viz[y:y+20, legend_x:legend_x+20] = color

        return enhanced_viz