"""
ComfyUI-GROUT: Custom nodes for GROUT mosaic grout line segmentation
"""

import os
import torch
import numpy as np
from PIL import Image

# Try to import ComfyUI utilities
try:
    import folder_paths
except ImportError:
    folder_paths = None


class GROUTSegmentation:
    """
    GROUT Segmentation Node
    Segments grout lines in mosaic images using the GROUT U-Net model.
    Outputs both probability heatmap and binary mask.
    """

    def __init__(self):
        self.model = None
        self.current_model_path = None
        self.current_device = None

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "device": (["cuda", "cpu"],),
                "mode": (["scale", "sliding_window"],),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider"
                }),
                "patch_size": ("INT", {
                    "default": 512,
                    "min": 128,
                    "max": 1024,
                    "step": 64,
                    "tooltip": "Size of patches for sliding window mode"
                }),
                "overlap": ("INT", {
                    "default": 64,
                    "min": 0,
                    "max": 256,
                    "step": 16,
                    "tooltip": "Overlap between patches in sliding window mode"
                }),
            },
            "optional": {
                "model_path": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Leave empty to auto-download from HuggingFace"
                }),
            }
        }

    RETURN_TYPES = ("IMAGE", "IMAGE", "MASK")
    RETURN_NAMES = ("probability_heatmap", "binary_mask", "mask")
    FUNCTION = "segment"
    CATEGORY = "GROUT"

    def get_default_model_path(self):
        """Get default model path in ComfyUI models directory."""
        if folder_paths:
            models_dir = folder_paths.models_dir
        else:
            # Fallback for standalone testing
            models_dir = os.path.join(os.path.dirname(__file__), "models")

        grout_dir = os.path.join(models_dir, "grout")
        os.makedirs(grout_dir, exist_ok=True)
        return os.path.join(grout_dir, "grout_b3_zeroshot_v1.pth")

    def download_model(self, target_path):
        """Download GROUT model from HuggingFace."""
        print("[GROUT] Downloading model from HuggingFace...")
        try:
            from huggingface_hub import hf_hub_download
            hf_hub_download(
                repo_id="advancedtech-sk/GROUT",
                filename="grout_b3_zeroshot_v1.pth",
                local_dir=os.path.dirname(target_path),
                local_dir_use_symlinks=False
            )
            print(f"[GROUT] Model downloaded to: {target_path}")
            return True
        except ImportError:
            print("[GROUT] ERROR: huggingface_hub not installed. Run: pip install huggingface_hub")
            return False
        except Exception as e:
            print(f"[GROUT] ERROR downloading model: {e}")
            return False

    def load_model(self, model_path, device):
        """Load or reload the GROUT model."""
        # Check if we need to reload
        if (self.model is not None and
            self.current_model_path == model_path and
            self.current_device == device):
            return self.model

        # Determine model path
        if not model_path or model_path.strip() == "":
            model_path = self.get_default_model_path()

        # Download if not exists
        if not os.path.exists(model_path):
            if not self.download_model(model_path):
                raise FileNotFoundError(f"Model not found and download failed: {model_path}")

        # Check CUDA availability
        if device == "cuda" and not torch.cuda.is_available():
            print("[GROUT] CUDA not available, falling back to CPU")
            device = "cpu"

        print(f"[GROUT] Loading model from: {model_path}")
        print(f"[GROUT] Device: {device}")

        # Import segmentation_models_pytorch
        try:
            import segmentation_models_pytorch as smp
        except ImportError:
            raise ImportError("segmentation_models_pytorch not installed. Run: pip install segmentation-models-pytorch")

        # Create model
        model = smp.Unet(encoder_name="efficientnet-b3", in_channels=3, classes=1)

        # Load weights
        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        else:
            model.load_state_dict(checkpoint)

        model.to(device)
        model.eval()

        self.model = model
        self.current_model_path = model_path
        self.current_device = device

        print("[GROUT] Model loaded successfully")
        return model

    def segment(self, image, device, mode, threshold, patch_size, overlap, model_path=""):
        """
        Segment grout lines in the input image.

        Args:
            image: ComfyUI IMAGE tensor [B, H, W, C] in range [0, 1]
            device: "cuda" or "cpu"
            mode: "scale" or "sliding_window"
            threshold: Threshold for binary mask (0.0 - 1.0)
            patch_size: Size of patches for sliding window mode
            overlap: Overlap between patches in sliding window mode
            model_path: Optional custom model path

        Returns:
            probability_heatmap: Grayscale probability map as IMAGE
            binary_mask: Thresholded mask as IMAGE
            mask: ComfyUI MASK format
        """
        # Load model
        model = self.load_model(model_path, device)

        # Process batch
        batch_size = image.shape[0]
        prob_maps = []
        binary_masks = []
        masks = []

        # ImageNet normalization
        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3)
        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3)

        for i in range(batch_size):
            # Get single image [H, W, C]
            img = image[i]
            original_h, original_w = img.shape[0], img.shape[1]

            if mode == "scale":
                prob_map = self._inference_scale(model, img, mean, std, device, original_h, original_w)
            else:
                prob_map = self._inference_sliding_window(model, img, mean, std, device, original_h, original_w, patch_size, overlap)

            # Create probability heatmap as RGB (grayscale) [H, W, 3]
            prob_rgb = prob_map.unsqueeze(-1).repeat(1, 1, 3)
            prob_maps.append(prob_rgb)

            # Create binary mask
            binary = (prob_map > threshold).float()
            binary_rgb = binary.unsqueeze(-1).repeat(1, 1, 3)
            binary_masks.append(binary_rgb)

            # MASK format [H, W]
            masks.append(binary)

        # Stack batches
        prob_heatmap = torch.stack(prob_maps, dim=0)
        binary_mask_img = torch.stack(binary_masks, dim=0)
        mask_out = torch.stack(masks, dim=0)

        return (prob_heatmap, binary_mask_img, mask_out)

    def _inference_scale(self, model, img, mean, std, device, original_h, original_w):
        """Run inference by scaling image to 512x512."""
        # Normalize
        img_norm = (img - mean) / std

        # Convert to [C, H, W] and add batch dim
        img_tensor = img_norm.permute(2, 0, 1).unsqueeze(0).to(device)

        # Resize to 512x512 for model
        img_resized = torch.nn.functional.interpolate(
            img_tensor, size=(512, 512), mode='bilinear', align_corners=False
        )

        # Inference
        with torch.no_grad():
            pred = model(img_resized.float())
            pred = torch.sigmoid(pred)

        # Resize back to original size
        pred_resized = torch.nn.functional.interpolate(
            pred, size=(original_h, original_w), mode='bilinear', align_corners=False
        )

        return pred_resized.squeeze().cpu()

    def _inference_sliding_window(self, model, img, mean, std, device, original_h, original_w, patch_size, overlap):
        """Run inference using sliding window for large images."""
        # If image is smaller than patch size, fall back to scale mode
        if original_h <= patch_size and original_w <= patch_size:
            return self._inference_scale(model, img, mean, std, device, original_h, original_w)

        # Create accumulation tensors
        prob_map = torch.zeros((original_h, original_w), dtype=torch.float32)
        count_map = torch.zeros((original_h, original_w), dtype=torch.float32)

        stride = patch_size - overlap

        # Calculate patch positions
        y_positions = list(range(0, original_h - patch_size + 1, stride))
        if original_h > patch_size and (not y_positions or y_positions[-1] + patch_size < original_h):
            y_positions.append(original_h - patch_size)
        if not y_positions:
            y_positions = [0]

        x_positions = list(range(0, original_w - patch_size + 1, stride))
        if original_w > patch_size and (not x_positions or x_positions[-1] + patch_size < original_w):
            x_positions.append(original_w - patch_size)
        if not x_positions:
            x_positions = [0]

        total_patches = len(y_positions) * len(x_positions)
        print(f"[GROUT] Processing {total_patches} patches ({len(y_positions)}x{len(x_positions)})...")

        for y in y_positions:
            for x in x_positions:
                # Extract patch [H, W, C]
                patch = img[y:y+patch_size, x:x+patch_size]

                # Normalize
                patch_norm = (patch - mean) / std

                # Convert to [C, H, W] and add batch dim
                patch_tensor = patch_norm.permute(2, 0, 1).unsqueeze(0).to(device)

                # Inference (patch is already patch_size x patch_size)
                with torch.no_grad():
                    pred = model(patch_tensor.float())
                    pred = torch.sigmoid(pred)

                # Get probability [H, W]
                patch_prob = pred.squeeze().cpu()

                # Accumulate
                prob_map[y:y+patch_size, x:x+patch_size] += patch_prob
                count_map[y:y+patch_size, x:x+patch_size] += 1

        # Average overlapping regions
        prob_map = prob_map / count_map.clamp(min=1)

        return prob_map


class GROUTThreshold:
    """
    GROUT Threshold Node
    Takes a probability heatmap and applies thresholding to create binary mask.
    Allows tuning threshold without re-running inference.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "probability_heatmap": ("IMAGE",),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "slider"
                }),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("binary_mask", "mask")
    FUNCTION = "apply_threshold"
    CATEGORY = "GROUT"

    def apply_threshold(self, probability_heatmap, threshold):
        """
        Apply threshold to probability heatmap.

        Args:
            probability_heatmap: Grayscale probability map as IMAGE [B, H, W, C]
            threshold: Threshold value (0.0 - 1.0)

        Returns:
            binary_mask: Thresholded mask as IMAGE
            mask: ComfyUI MASK format
        """
        batch_size = probability_heatmap.shape[0]
        binary_masks = []
        masks = []

        for i in range(batch_size):
            # Get probability (use first channel, should be grayscale)
            prob = probability_heatmap[i, :, :, 0]

            # Apply threshold
            binary = (prob > threshold).float()

            # Create RGB binary mask
            binary_rgb = binary.unsqueeze(-1).repeat(1, 1, 3)
            binary_masks.append(binary_rgb)

            # MASK format
            masks.append(binary)

        binary_mask_img = torch.stack(binary_masks, dim=0)
        mask_out = torch.stack(masks, dim=0)

        return (binary_mask_img, mask_out)


class GROUTHeatmapColorize:
    """
    GROUT Heatmap Colorize Node
    Converts grayscale probability heatmap to colored visualization.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "probability_heatmap": ("IMAGE",),
                "colormap": (["magma", "viridis", "plasma", "inferno", "hot", "jet"],),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("colored_heatmap",)
    FUNCTION = "colorize"
    CATEGORY = "GROUT"

    def colorize(self, probability_heatmap, colormap):
        """
        Apply colormap to probability heatmap.
        """
        import matplotlib.pyplot as plt

        batch_size = probability_heatmap.shape[0]
        colored_maps = []

        cmap = plt.get_cmap(colormap)

        for i in range(batch_size):
            # Get probability (grayscale)
            prob = probability_heatmap[i, :, :, 0].numpy()

            # Apply colormap
            colored = cmap(prob)[:, :, :3]  # Remove alpha channel

            colored_maps.append(torch.from_numpy(colored).float())

        colored_output = torch.stack(colored_maps, dim=0)

        return (colored_output,)


# Node mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "GROUTSegmentation": GROUTSegmentation,
    "GROUTThreshold": GROUTThreshold,
    "GROUTHeatmapColorize": GROUTHeatmapColorize,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "GROUTSegmentation": "GROUT Segmentation",
    "GROUTThreshold": "GROUT Threshold",
    "GROUTHeatmapColorize": "GROUT Heatmap Colorize",
}