"""
Map Utilities
Height and AO map processing
"""

import torch


class LotusHeightProcessor:
    """
    Process Lotus depth/height map output
    Connect: LotusSampler → VAEDecode → This node
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "lotus_depth": ("IMAGE",),
                "invert": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Invert height values"
                }),
                "bit_depth": (["8-bit", "16-bit", "32-bit"], {
                    "default": "16-bit",
                    "tooltip": "Output bit depth (save as EXR for 16/32-bit)"
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("height",)
    FUNCTION = "process"
    CATEGORY = "Texture Alchemist/Maps"
    
    def process(self, lotus_depth, invert, bit_depth):
        """Process height map"""
        
        height = lotus_depth.clone()
        
        # Convert to grayscale
        if height.shape[-1] == 3:
            weights = torch.tensor([0.299, 0.587, 0.114], 
                                   device=height.device, dtype=height.dtype)
            height = torch.sum(height * weights, dim=-1, keepdim=True)
            height = height.repeat(1, 1, 1, 3)
        
        # Normalize to 0-1
        h_min = height.min()
        h_max = height.max()
        if h_max - h_min > 1e-6:
            height = (height - h_min) / (h_max - h_min)
        
        # Invert if requested
        if invert:
            height = 1.0 - height
        
        print(f"✓ Height processed ({bit_depth})")
        if bit_depth in ["16-bit", "32-bit"]:
            print(f"  Note: Save as OpenEXR to preserve {bit_depth} precision")
        
        return (height,)


class AOApproximator:
    """
    Generate Ambient Occlusion map from height and/or normal maps
    - With height: Uses height sampling to detect concave areas that would be occluded
    - Without height: Uses normal map orientation to approximate basic AO
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "radius": ("INT", {
                    "default": 8,
                    "min": 1,
                    "max": 64,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Sampling radius in pixels (larger = broader occlusion)"
                }),
                "strength": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "AO intensity multiplier"
                }),
                "samples": ("INT", {
                    "default": 16,
                    "min": 4,
                    "max": 64,
                    "step": 1,
                    "display": "number",
                    "tooltip": "Number of sampling directions (more = better quality, slower)"
                }),
                "contrast": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.0,
                    "max": 3.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "Contrast adjustment for AO"
                }),
                "use_normal": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Use normal map to improve AO quality (if provided)"
                }),
            },
            "optional": {
                "height": ("IMAGE",),
                "normal": ("IMAGE",),
            },
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("ao",)
    FUNCTION = "generate_ao"
    CATEGORY = "Texture Alchemist/Maps"
    
    def generate_ao(self, radius, strength, samples, contrast, use_normal, height=None, normal=None):
        """Generate AO from height and/or normal map"""
        
        print("\n" + "="*60)
        print("AO Approximator")
        print("="*60)
        
        # Check inputs
        if height is None and normal is None:
            print("⚠ Warning: No height or normal map provided, returning neutral AO")
            # Create a neutral gray AO map (50% occlusion)
            ao_rgb = torch.ones((1, 512, 512, 3), dtype=torch.float32) * 0.7
            print("="*60 + "\n")
            return (ao_rgb,)
        
        if height is not None:
            print(f"Height shape: {height.shape}")
        if normal is not None:
            print(f"Normal shape: {normal.shape}")
        print(f"Radius: {radius}, Samples: {samples}, Strength: {strength}, Contrast: {contrast}")
        print(f"Use normal: {use_normal}")
        
        # Determine reference shape and device
        if height is not None:
            reference = height
            batch, h, w, channels = height.shape
        else:
            reference = normal
            batch, h, w, channels = normal.shape
        
        # Convert height to grayscale if needed
        if height is not None:
            height_map = height.clone()
            if height_map.shape[-1] == 3:
                weights = torch.tensor([0.299, 0.587, 0.114], 
                                       device=height_map.device, dtype=height_map.dtype)
                height_map = torch.sum(height_map * weights, dim=-1, keepdim=True)
            elif height_map.shape[-1] != 1:
                height_map = height_map[..., 0:1]
        else:
            height_map = None
        
        # Initialize AO map (start fully lit)
        ao = torch.ones((batch, h, w, 1), device=reference.device, dtype=reference.dtype)
        
        # If we have a height map, use height-based AO
        if height_map is not None:
            # Generate sampling angles
            import math
            angles = [2.0 * math.pi * i / samples for i in range(samples)]
            
            print(f"\n⚙ Computing AO with {samples} samples at radius {radius}...")
            
            # For each sampling direction
            occlusion_sum = torch.zeros_like(ao)
            
            for angle in angles:
                # Calculate offset
                dx = int(round(math.cos(angle) * radius))
                dy = int(round(math.sin(angle) * radius))
                
                # Skip if no offset
                if dx == 0 and dy == 0:
                    continue
                
                # Sample height at offset position (with boundary handling)
                # Pad the height map to handle boundaries
                pad_h = abs(dy) if dy != 0 else 0
                pad_w = abs(dx) if dx != 0 else 0
                
                # Use reflection padding
                height_padded = torch.nn.functional.pad(
                    height_map.permute(0, 3, 1, 2),  # BHWC -> BCHW
                    (pad_w, pad_w, pad_h, pad_h),
                    mode='replicate'
                ).permute(0, 2, 3, 1)  # BCHW -> BHWC
                
                # Calculate the sampling positions
                y_start = max(0, dy) + pad_h
                y_end = y_start + h
                x_start = max(0, dx) + pad_w
                x_end = x_start + w
                
                # Center position
                center_y = pad_h
                center_x = pad_w
                
                # Sample neighbor and center
                neighbor_height = height_padded[:, y_start:y_end, x_start:x_end, :]
                center_height = height_padded[:, center_y:center_y+h, center_x:center_x+w, :]
                
                # Calculate height difference (positive if neighbor is higher = occlusion)
                height_diff = neighbor_height - center_height
                
                # Convert to occlusion (only positive differences contribute)
                occlusion = torch.clamp(height_diff, 0.0, 1.0)
                
                occlusion_sum += occlusion
            
            # Average the occlusion
            if samples > 0:
                occlusion_avg = occlusion_sum / samples
            else:
                occlusion_avg = occlusion_sum
        else:
            # No height map - generate basic AO from normal orientation if available
            print("\n⚙ No height map - generating basic AO from normal orientation...")
            occlusion_avg = torch.zeros((batch, h, w, 1), device=reference.device, dtype=reference.dtype)
        
        # Apply strength
        occlusion_avg = occlusion_avg * strength
        
        # Convert to AO (1.0 = no occlusion, 0.0 = full occlusion)
        ao = 1.0 - torch.clamp(occlusion_avg, 0.0, 1.0)
        
        # If we don't have height but have normal, generate AO from normal map
        if height_map is None and normal is not None:
            print("✓ Generating AO from normal map orientation")
            
            if normal.shape[-1] >= 3:
                # Use blue channel (Z/up direction) and curvature
                normal_z = normal[..., 2:3]
                
                # Resize to match AO if needed
                if normal_z.shape[1:3] != ao.shape[1:3]:
                    normal_z = torch.nn.functional.interpolate(
                        normal_z.permute(0, 3, 1, 2),
                        size=ao.shape[1:3],
                        mode='bilinear',
                        align_corners=False
                    ).permute(0, 2, 3, 1)
                
                # Generate AO based on surface orientation
                # Faces pointing up (normal_z ~1.0) = bright (little AO)
                # Faces pointing sideways/down (normal_z ~0.5 or less) = dark (more AO)
                # Map range: 1.0 (up) -> 0.9 AO, 0.5 (sideways) -> 0.6 AO, 0.0 (down) -> 0.3 AO
                ao = normal_z * 0.6 + 0.3
                ao = torch.clamp(ao, 0.0, 1.0)
        
        # Apply contrast
        if contrast != 1.0:
            ao = (ao - 0.5) * contrast + 0.5
            ao = torch.clamp(ao, 0.0, 1.0)
        
        # If we have both height-based AO and normal, use normal to enhance/bias
        if height_map is not None and use_normal and normal is not None:
            print("✓ Applying normal-based bias to height AO")
            
            # Convert normal to grayscale or use blue channel (up direction)
            if normal.shape[-1] >= 3:
                # Use blue channel (Z/up direction)
                # Areas facing up get less AO, areas facing down/sideways get more
                normal_z = normal[..., 2:3]
                
                # Resize to match AO if needed
                if normal_z.shape[1:3] != ao.shape[1:3]:
                    normal_z = torch.nn.functional.interpolate(
                        normal_z.permute(0, 3, 1, 2),
                        size=ao.shape[1:3],
                        mode='bilinear',
                        align_corners=False
                    ).permute(0, 2, 3, 1)
                
                # Blend AO with normal bias
                # Faces pointing up (normal_z close to 1.0) = less AO
                # Faces pointing sideways/down (normal_z close to 0.5 or less) = more AO
                normal_bias = torch.clamp((1.0 - normal_z) * 0.5, 0.0, 0.5)
                ao = ao * (1.0 - normal_bias)
        
        # Convert to RGB
        ao_rgb = ao.repeat(1, 1, 1, 3)
        
        print(f"\n✓ AO generated")
        print(f"  Range: [{ao_rgb.min():.3f}, {ao_rgb.max():.3f}]")
        print("\n" + "="*60)
        print("✓ AO Approximation Complete")
        print("="*60 + "\n")
        
        return (ao_rgb,)


class GammaAdjust:
    """
    Apply gamma correction to images
    Gamma < 1.0 brightens, Gamma > 1.0 darkens
    Common values: 0.45 (linear to sRGB), 2.2 (sRGB to linear)
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "gamma": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 5.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "Gamma value (< 1.0 = brighten, > 1.0 = darken, 1.0 = no change)"
                }),
                "per_channel": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Apply gamma to each channel independently (useful for color correction)"
                }),
            },
            "optional": {
                "gamma_red": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 5.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "Gamma for red channel (only used if per_channel is enabled)"
                }),
                "gamma_green": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 5.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "Gamma for green channel (only used if per_channel is enabled)"
                }),
                "gamma_blue": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 5.0,
                    "step": 0.01,
                    "display": "number",
                    "tooltip": "Gamma for blue channel (only used if per_channel is enabled)"
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "adjust_gamma"
    CATEGORY = "Texture Alchemist/Adjustment"
    
    def adjust_gamma(self, image, gamma, per_channel, gamma_red=1.0, gamma_green=1.0, gamma_blue=1.0):
        """Apply gamma correction to image"""
        
        print("\n" + "="*60)
        print("Gamma Adjust")
        print("="*60)
        print(f"Input shape: {image.shape}")
        print(f"Gamma: {gamma}")
        print(f"Per channel: {per_channel}")
        
        # Check if we need to do anything
        if not per_channel and gamma == 1.0:
            print("✓ Gamma is 1.0, no adjustment needed")
            print("="*60 + "\n")
            return (image,)
        
        if per_channel and gamma_red == 1.0 and gamma_green == 1.0 and gamma_blue == 1.0:
            print("✓ All channel gammas are 1.0, no adjustment needed")
            print("="*60 + "\n")
            return (image,)
        
        # Clamp input to valid range [0, 1]
        image_clamped = torch.clamp(image, 0.0, 1.0)
        
        if per_channel and image.shape[-1] >= 3:
            # Apply different gamma to each channel
            print(f"  Red gamma: {gamma_red}")
            print(f"  Green gamma: {gamma_green}")
            print(f"  Blue gamma: {gamma_blue}")
            
            # Split channels
            r = image_clamped[..., 0:1]
            g = image_clamped[..., 1:2]
            b = image_clamped[..., 2:3]
            
            # Apply gamma to each channel
            r_adjusted = torch.pow(r, gamma_red)
            g_adjusted = torch.pow(g, gamma_green)
            b_adjusted = torch.pow(b, gamma_blue)
            
            # Recombine
            result = torch.cat([r_adjusted, g_adjusted, b_adjusted], dim=-1)
            
            # If there was an alpha channel, preserve it
            if image.shape[-1] == 4:
                alpha = image_clamped[..., 3:4]
                result = torch.cat([result, alpha], dim=-1)
                print("  ✓ Alpha channel preserved")
            elif image.shape[-1] > 4:
                # Preserve extra channels
                extra = image_clamped[..., 3:]
                result = torch.cat([result, extra], dim=-1)
        else:
            # Apply same gamma to all channels
            result = torch.pow(image_clamped, gamma)
        
        print(f"\n✓ Gamma correction applied")
        print(f"  Input range: [{image.min():.3f}, {image.max():.3f}]")
        print(f"  Output range: [{result.min():.3f}, {result.max():.3f}]")
        print("="*60 + "\n")
        
        return (result,)


class SimpleGammaAdjust:
    """
    Simple gamma correction with just one slider
    Gamma < 1.0 brightens, Gamma > 1.0 darkens
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "gamma": ("FLOAT", {
                    "default": 1.0,
                    "min": 0.1,
                    "max": 5.0,
                    "step": 0.01,
                    "tooltip": "Gamma value (< 1.0 = brighten, > 1.0 = darken, 1.0 = no change)"
                }),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "adjust_gamma"
    CATEGORY = "Texture Alchemist/Adjustment"
    
    def adjust_gamma(self, image, gamma):
        """Apply gamma correction to image"""
        # Clamp to valid range
        image_clamped = torch.clamp(image, 0.0, 1.0)
        
        # Apply gamma
        result = torch.pow(image_clamped, gamma)
        
        return (result,)


# Node registration
NODE_CLASS_MAPPINGS = {
    "LotusHeightProcessor": LotusHeightProcessor,
    "AOApproximator": AOApproximator,
    "GammaAdjust": GammaAdjust,
    "SimpleGammaAdjust": SimpleGammaAdjust,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LotusHeightProcessor": "Height Processor (Lotus)",
    "AOApproximator": "AO Approximator",
    "GammaAdjust": "Gamma Adjust",
    "SimpleGammaAdjust": "Simple Gamma Adjust",
}