"""ComfyUI Node for D-DFFNet Defocus Blur Detection.

This node performs defocus blur detection using the D-DFFNet model.
The model should be placed in <ComfyUI_root>/models/dffnet/ directory.
"""
import os
import numpy as np
import torch
import torch.nn.functional as F

# ComfyUI imports
import folder_paths

import comfy.model_management as model_management

from .model.PDNet_resnest import PDnet_ResNest_4F_SL_LA_D


# Register our custom model folder
folder_paths.folder_names_and_paths["dffnet"] = (
    [os.path.join(folder_paths.models_dir, "dffnet")],
    [".pth", ".pt", ".pth.tar"],
)


class DDFFNetNode:
    """ComfyUI node for D-DFFNet defocus blur detection."""

    CATEGORY = "D-DFFNet"
    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("mask",)
    FUNCTION = "detect"

    @classmethod
    def INPUT_TYPES(cls):
        model_names = cls._get_model_list()
        return {
            "required": {
                "image": ("IMAGE",),
                "model_name": (model_names,),
            },
            "optional": {
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "Threshold for binarization (0 = output raw probability)"
                }),
            },
        }

    @classmethod
    def _get_model_list(cls):
        """Get list of available model files from ComfyUI/models/dffnet/"""
        paths, _ = folder_paths.folder_names_and_paths.get("dffnet", ([], []))
        models = []
        for folder in paths:
            if os.path.isdir(folder):
                for f in os.listdir(folder):
                    if f.endswith((".pth", ".pt", ".pth.tar")):
                        models.append(f)
        return sorted(models) if models else ["[No models found]"]

    def detect(self, image, model_name, threshold=0.5):
        """
        Run D-DFFNet inference on the input image.

        Args:
            image: ComfyUI IMAGE tensor, shape (B, H, W, 3), float32, range [0, 1]
            model_name: filename of the model weight to load
            threshold: if > 0, binarize output at this threshold; if 0, output raw sigmoid probabilities
        Returns:
            mask: ComfyUI MASK tensor, shape (B, H, W), float32, range [0, 1]
        """
        device = self._get_device()

        # Load model
        model_path = self._find_model_path(model_name)
        if model_path is None:
            raise FileNotFoundError(f"Model '{model_name}' not found in ComfyUI/models/dffnet/")

        model = self._load_model(model_path, device)

        # Prepare input: (B, H, W, 3) -> (B, 3, H, W), normalize to [0,1] already
        B, H, W, C = image.shape
        input_tensor = image.permute(0, 3, 1, 2).to(device)  # BHWC -> BCHW

        # Preprocess: resize to 320x320 + ImageNet normalization
        input_320 = F.interpolate(input_tensor, size=(320, 320), mode="bilinear", align_corners=True)
        mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1)
        input_norm = (input_320 - mean) / std

        # Inference
        with torch.no_grad():
            pred, _, _, _ = model(input_norm)
            pred = torch.sigmoid(pred)  # (B, 1, 320, 320)

        # Post-process: resize back to original image size
        pred = F.interpolate(pred, size=(H, W), mode="bilinear", align_corners=True)  # (B, 1, H, W)

        # Apply threshold if specified
        if threshold > 0:
            pred = (pred > threshold).float()

        # Convert to MASK format: (B, 1, H, W) -> (B, H, W)
        mask = pred.squeeze(1)  # (B, H, W)

        # Move to CPU and ensure correct dtype
        mask = mask.cpu().float()

        return (mask,)

    def _get_device(self):
        return model_management.get_torch_device()

    def _find_model_path(self, model_name):
        """Find the full path to a model file in ComfyUI/models/dffnet/"""
        if model_name.startswith("[") and model_name.endswith("]"):
            return None
        paths, _ = folder_paths.folder_names_and_paths.get("dffnet", ([], []))
        for folder in paths:
            candidate = os.path.join(folder, model_name)
            if os.path.isfile(candidate):
                return candidate
        return None

    def _load_model(self, model_path, device):
        """Load D-DFFNet model from checkpoint."""
        model = PDnet_ResNest_4F_SL_LA_D()
        state_dict = torch.load(model_path, map_location=device)

        # Handle checkpoint formats: remove 'module.' prefix if present (DataParallel)
        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith("module."):
                new_state_dict[k[7:]] = v  # strip 'module.' prefix
            else:
                new_state_dict[k] = v

        # Load with strict=False: side outputs (side5-side8) are only used during training
        # and won't be present if loading a DFFNet checkpoint instead of D-DFFNet.
        # Missing depth side outputs (side5-side8) are normal when loading DFFNet weights;
        # they don't affect inference quality since only the main output path is used.
        missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
        if missing_keys:
            dffnet_keys = [k for k in missing_keys if 'side5' in k or 'side6' in k or 'side7' in k or 'side8' in k]
            other_keys = [k for k in missing_keys if k not in dffnet_keys]
            if dffnet_keys:
                print(f"[D-DFFNet] Missing depth side outputs (expected if loading DFFNet .pth): {dffnet_keys}")
            if other_keys:
                print(f"[D-DFFNet] Missing keys: {other_keys}")
        if unexpected_keys:
            print(f"[D-DFFNet] Unexpected keys: {unexpected_keys}")

        model.to(device)
        model.eval()
        return model