# ComfyUI Multi-Analysis Heatmaps (Extended & Optimized, Adaptive SSIM Visualization)
# Version: 3.0.3
# Parallel outputs:
#  - SSIM, DSSIM, MS-SSIM  (adaptive, zero-sensitive visualization via percentile capping)
#  - LPIPS (approx via VGG16 features, optional)
#  - Per-pixel abs diff, Residual (signed)
#  - FFT spectra (orig, upscaled, diff)
#  - ΔE2000 heatmap + L*, a*, b* diffs (Lab)
#  - Gradient diff (Scharr) + Edge IoU ("Canny-like")
#  - Wavelet band diff (Haar, multi-level)
#  - Denoiser residuum (Gaussian high-pass)
#  - Round-trip stability map (downscale->upscale)
#
# Perf:
#  - dtype-safe fp32, autocast off, no_grad
#  - cached Gaussian kernels & VGG16 features
#
# No external deps beyond torch/torchvision (optional for LPIPS).
# Author: chatgpt

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# ======================= Node Meta =======================
NODE_NAME = "Image Analysis (SSIM/LPIPS/FFT/ΔE2000)"
NODE_CATEGORY = "Image/Analysis"
NODE_AUTHOR = "chatgpt"
NODE_VERSION = "3.0.3"

# ======================= Caches ==========================
_GAUSS2D_CACHE = {}   # (dev, idx, dtype, ch, win, sigma) -> [C,1,win,win]
_GAUSS1D_CACHE = {}   # (dev, idx, dtype, win, sigma) -> [1,1,1,W]
_VGG_CACHE = {}       # (dev, idx) -> {"model": VGG16Features, "mean": t, "std": t}

def _dev_key(device: torch.device, dtype=None):
    return (device.type, device.index if device.index is not None else -1, str(dtype) if dtype else "")

# ======================= Tensor helpers ==================
def _squeeze_to_4d(img: torch.Tensor) -> torch.Tensor:
    """
    Squeeze superfluous singleton dims (==1) until rank <= 4.
    Keeps batch/channels semantics; only squeezes when rank > 4.
    """
    if not isinstance(img, torch.Tensor):
        raise TypeError(f"_squeeze_to_4d expected torch.Tensor, got {type(img)}")
    t = img
    while t.dim() > 4 and 1 in t.shape:
        for d, s in enumerate(t.shape):
            if s == 1:
                t = t.squeeze(d)
                break
    return t

def to_bchw(img: torch.Tensor) -> torch.Tensor:
    """
    Robust conversion to [B,C,H,W].
    Accepts [B,H,W,C], [B,C,H,W], [H,W,C], [C,H,W], [H,W] and higher ranks with singleton dims.
    """
    if not isinstance(img, torch.Tensor):
        raise TypeError(f"to_bchw expected torch.Tensor, got {type(img)}")
    x = _squeeze_to_4d(img)

    if x.dim() == 4:
        if x.shape[-1] in (1, 3):  # BHWC -> BCHW
            return x.permute(0, 3, 1, 2).contiguous()
        return x.contiguous()       # assume BCHW

    if x.dim() == 3:
        if x.shape[-1] in (1, 3):   # [H,W,C] -> [1,C,H,W]
            return x.permute(2, 0, 1).unsqueeze(0).contiguous()
        if x.shape[0] in (1, 3):    # [C,H,W] -> [1,C,H,W]
            return x.unsqueeze(0).contiguous()
        return x.unsqueeze(0).contiguous()  # fallback assume [C,H,W]

    if x.dim() == 2:
        return x.unsqueeze(0).unsqueeze(0).contiguous()  # [H,W] -> [1,1,H,W]

    raise ValueError(f"to_bchw expected 2D/3D/4D tensor after squeeze, got shape {tuple(x.shape)}")

def to_bhwc(img: torch.Tensor) -> torch.Tensor:
    """
    Robust conversion to [B,H,W,C].
    Accepts [B,C,H,W], [B,H,W,C], [C,H,W], [H,W,C], [H,W] and higher ranks with singleton dims.
    """
    if not isinstance(img, torch.Tensor):
        raise TypeError(f"to_bhwc expected torch.Tensor, got {type(img)}")
    x = _squeeze_to_4d(img)

    if x.dim() == 4:
        if x.shape[-1] in (1, 3):   # already BHWC
            return x.contiguous()
        return x.permute(0, 2, 3, 1).contiguous()  # BCHW -> BHWC

    if x.dim() == 3:
        if x.shape[-1] in (1, 3):   # [H,W,C] -> [1,H,W,C]
            return x.unsqueeze(0).contiguous()
        if x.shape[0] in (1, 3):    # [C,H,W] -> [1,H,W,C]
            return x.permute(1, 2, 0).unsqueeze(0).contiguous()
        return x.permute(1, 2, 0).unsqueeze(0).contiguous()  # fallback assume [C,H,W]

    if x.dim() == 2:
        return x.unsqueeze(0).unsqueeze(-1).contiguous()     # [H,W] -> [1,H,W,1]

    raise ValueError(f"to_bhwc expected 2D/3D/4D tensor after squeeze, got shape {tuple(x.shape)}")

def _ensure_bhwc3(x: torch.Tensor) -> torch.Tensor:
    """
    Ensure BHWC with 3 channels for ComfyUI preview stability.
    """
    t = to_bhwc(x)
    if t.shape[-1] == 1:
        t = t.repeat(1, 1, 1, 3)
    return t

def ensure_same_device(*tensors):
    dev = None
    for t in tensors:
        if isinstance(t, torch.Tensor):
            dev = t.device
            break
    if dev is None:
        dev = torch.device("cpu")
    out = []
    for t in tensors:
        out.append(t.to(dev) if isinstance(t, torch.Tensor) else t)
    return out

def resize_like(img_bchw: torch.Tensor, target_bchw: torch.Tensor) -> torch.Tensor:
    _, _, H, W = target_bchw.shape
    return F.interpolate(img_bchw, size=(H, W), mode="bilinear", align_corners=False)

def center_crop_to_min(a: torch.Tensor, b: torch.Tensor):
    _, _, Ha, Wa = a.shape
    _, _, Hb, Wb = b.shape
    H = min(Ha, Hb); W = min(Wa, Wb)
    def crop(x):
        _, _, Hx, Wx = x.shape
        top = (Hx - H) // 2
        left = (Wx - W) // 2
        return x[:, :, top:top+H, left:left+W]
    return crop(a), crop(b)

def pad_to_max(a: torch.Tensor, b: torch.Tensor):
    _, _, Ha, Wa = a.shape
    _, _, Hb, Wb = b.shape
    H = max(Ha, Hb); W = max(Wa, Wb)
    def pad(x):
        _, _, Hx, Wx = x.shape
        pt = (H - Hx) // 2
        pb = H - Hx - pt
        pl = (W - Wx) // 2
        pr = W - Wx - pl
        return F.pad(x, (pl, pr, pt, pb), mode="constant", value=0.0)
    return pad(a), pad(b)

# ======================= Normalization & Colormaps =======
def _ramp(x, a, b):
    return torch.clamp((x - a) / (b - a), 0.0, 1.0)

def jet_colormap(v01: torch.Tensor) -> torch.Tensor:
    r = _ramp(v01, 0.35, 0.66) + _ramp(v01, 0.66, 1.0)
    g = _ramp(v01, 0.0, 0.125) + _ramp(v01, 0.125, 0.375) + _ramp(v01, 0.375, 0.66) - _ramp(v01, 0.66, 0.875)
    b = _ramp(v01, 0.0, 0.375) + _ramp(v01, 0.375, 0.66)
    return torch.clamp(torch.cat([r, g, b], dim=1), 0.0, 1.0)

def grayscale_colormap(v01: torch.Tensor) -> torch.Tensor:
    return v01.repeat(1, 3, 1, 1)

def apply_colormap(v01: torch.Tensor, mode: str = "jet") -> torch.Tensor:
    return grayscale_colormap(v01) if mode == "grayscale" else jet_colormap(v01)

def normalize01(x: torch.Tensor) -> torch.Tensor:
    xmin = x.amin(dim=(2,3), keepdim=True)
    xmax = x.amax(dim=(2,3), keepdim=True)
    span = torch.clamp(xmax - xmin, min=1e-8)
    return torch.clamp((x - xmin) / span, 0.0, 1.0)

# ---- Fixed/Adaptive normalization helpers for difference-like maps (SSIM family)
def _percentile(x: torch.Tensor, q: float):
    """Returns per-sample percentile (q in [0,1]) for flattened spatial dims."""
    q = float(max(0.0, min(1.0, q)))
    flat = x.flatten(1)  # [B,HW]
    try:
        p = torch.quantile(flat, q, dim=1)  # [B]
    except Exception:
        vals, _ = torch.sort(flat, dim=1)
        k = int(q * (vals.shape[1] - 1))
        p = vals[:, k]
    return p.view(-1, 1, 1, 1)  # [B,1,1,1]

def normalize_heatmap_adaptive(x: torch.Tensor, base_min: float = 0.02, qcap: float = 0.98) -> torch.Tensor:
    """
    Map a non-negative difference map to [0,1] using an adaptive cap:
      cap = max(base_min, percentile(x, qcap)) per image.
    This keeps identical images black, avoids 'all white' saturation for small but widespread diffs.
    """
    cap = torch.clamp(_percentile(x, qcap), min=base_min)
    return torch.clamp(x / cap, 0.0, 1.0)

# ======================= Gaussian Kernels =================
def gaussian2d(window_size=11, sigma=1.5, channels=3, device="cpu", dtype=torch.float32):
    key = (_dev_key(torch.device(device), dtype), channels, int(window_size), float(sigma))
    k = _GAUSS2D_CACHE.get(key)
    if k is not None and k.device == torch.device(device) and k.dtype == dtype:
        return k
    coords = torch.arange(window_size, dtype=dtype, device=device) - (window_size - 1) / 2.0
    g = torch.exp(-(coords**2) / (2 * sigma * sigma))
    g = g / (g.sum() + 1e-12)
    k2d = torch.outer(g, g)
    k = k2d.expand(channels, 1, window_size, window_size).contiguous()
    _GAUSS2D_CACHE[key] = k
    return k

def gaussian1d(window_size=11, sigma=1.5, device="cpu", dtype=torch.float32):
    key = (_dev_key(torch.device(device), dtype), int(window_size), float(sigma))
    k = _GAUSS1D_CACHE.get(key)
    if k is not None and k.device == torch.device(device) and k.dtype == dtype:
        return k
    coords = torch.arange(window_size, dtype=dtype, device=device) - (window_size - 1) / 2.0
    g = torch.exp(-(coords**2) / (2 * sigma * sigma))
    g = g / (g.sum() + 1e-12)
    k = g.view(1, 1, 1, -1).contiguous()  # [1,1,1,W]
    _GAUSS1D_CACHE[key] = k
    return k

def gaussian_blur(img_bchw: torch.Tensor, win=11, sigma=1.5):
    # separable blur (faster), per-channel depthwise
    B, C, H, W = img_bchw.shape
    device, dtype = img_bchw.device, img_bchw.dtype
    k1 = gaussian1d(win, sigma, device, dtype)       # [1,1,1,W]
    k1 = k1.expand(C, 1, 1, -1).contiguous()         # [C,1,1,W]
    k1t = k1.transpose(2, 3).contiguous()            # [C,1,W,1]
    # horizontal
    x = F.conv2d(img_bchw, k1, groups=C, padding=(0, win//2))
    # vertical
    x = F.conv2d(x, k1t, groups=C, padding=(win//2, 0))
    return x

# ======================= SSIM / MS-SSIM ==================
def ssim_map(img1: torch.Tensor, img2: torch.Tensor, window_size=11, sigma=1.5, eps=1e-12):
    assert img1.shape == img2.shape, "SSIM input shapes must match"
    B, C, H, W = img1.shape
    device = img1.device
    dtype = img1.dtype
    C1 = torch.as_tensor((0.01 ** 2), device=device, dtype=dtype)
    C2 = torch.as_tensor((0.03 ** 2), device=device, dtype=dtype)
    win = gaussian2d(window_size, sigma, channels=C, device=device, dtype=dtype)

    mu1 = F.conv2d(img1, win, groups=C, padding=window_size//2)
    mu2 = F.conv2d(img2, win, groups=C, padding=window_size//2)

    mu1_sq = mu1 * mu1
    mu2_sq = mu2 * mu2
    mu1_mu2 = mu1 * mu2

    sigma1_sq = F.conv2d(img1 * img1, win, groups=C, padding=window_size//2) - mu1_sq
    sigma2_sq = F.conv2d(img2 * img2, win, groups=C, padding=window_size//2) - mu2_sq
    sigma12   = F.conv2d(img1 * img2, win, groups=C, padding=window_size//2) - mu1_mu2

    num = (2 * mu1_mu2 + C1) * (2 * sigma12 + C2)
    den = (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
    ssim_c = num / (den + eps)      # [B,C,H,W]
    ssim_mean = ssim_c.mean(dim=1, keepdim=True)
    return torch.clamp(ssim_mean, 0.0, 1.0)

def ms_ssim_map(img1: torch.Tensor, img2: torch.Tensor, scales=3):
    assert scales >= 1
    maps = []
    a = img1
    b = img2
    base_size = (img1.shape[2], img1.shape[3])
    for s in range(scales):
        m = ssim_map(a, b)
        if s > 0:
            m = F.interpolate(m, size=base_size, mode="bilinear", align_corners=False)
        maps.append(m)
        if s < scales - 1:
            a = F.avg_pool2d(a, kernel_size=2, stride=2)
            b = F.avg_pool2d(b, kernel_size=2, stride=2)
    return torch.clamp(torch.stack(maps, dim=0).mean(dim=0), 0.0, 1.0)

# ======================= LPIPS-lite (optional) ==========
class VGG16Features(nn.Module):
    def __init__(self):
        super().__init__()
        from torchvision import models
        vgg = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_FEATURES)
        feats = vgg.features
        self.b1 = nn.Sequential(*feats[:4])    # conv1_2
        self.b2 = nn.Sequential(*feats[4:9])   # conv2_2
        self.b3 = nn.Sequential(*feats[9:16])  # conv3_3
        self.b4 = nn.Sequential(*feats[16:23]) # conv4_3
        for p in self.parameters():
            p.requires_grad_(False)

def _vgg_get(device: torch.device, dtype: torch.dtype):
    key = (device.type, device.index if device.index is not None else -1)
    entry = _VGG_CACHE.get(key)
    if entry is not None:
        return entry["model"], entry["mean"].to(device=device, dtype=dtype), entry["std"].to(device=device, dtype=dtype)
    model = VGG16Features().to(device).eval()
    mean = torch.tensor([0.485, 0.456, 0.406], device=device, dtype=dtype).view(1,3,1,1)
    std  = torch.tensor([0.229, 0.224, 0.225], device=device, dtype=dtype).view(1,3,1,1)
    _VGG_CACHE[key] = {"model": model, "mean": mean, "std": std}
    return model, mean, std

def _imagenet_norm(x: torch.Tensor, mean: torch.Tensor, std: torch.Tensor):
    return (x.clamp(0,1) - mean) / std

def lpips_heatmap(a: torch.Tensor, b: torch.Tensor):
    device, dtype = a.device, a.dtype
    try:
        model, mean, std = _vgg_get(device, dtype)
    except Exception as e:
        return None, f"LPIPS disabled: {str(e)}"
    with torch.no_grad():
        A = _imagenet_norm(a, mean, std)
        B = _imagenet_norm(b, mean, std)
        a1 = model.b1(A); b1 = model.b1(B)
        a2 = model.b2(a1); b2 = model.b2(b1)
        a3 = model.b3(a2); b3 = model.b3(b2)
        a4 = model.b4(a3); b4 = model.b4(b3)
        feats_a = [a1, a2, a3, a4]
        feats_b = [b1, b2, b3, b4]
        maps = []
        H, W = a.shape[2], a.shape[3]
        for fa, fb in zip(feats_a, feats_b):
            d = (fa - fb).pow(2).mean(dim=1, keepdim=True)
            d = F.interpolate(d, size=(H, W), mode="bilinear", align_corners=False)
            maps.append(d)
        m = torch.stack(maps, dim=0).mean(dim=0)  # [B,1,H,W]
        m01 = normalize01(m)  # NOTE: LPIPS keeps per-image norm (heuristic visual)
        score = float(m.mean().item())
        return m01, f"LPIPS approx (VGG16 map mean): {score:.6f}"

# ======================= Color: sRGB <-> Lab (D65) =======
def _srgb_to_linear(x):
    a = 0.055
    return torch.where(x <= 0.04045, x/12.92, ((x + a)/(1 + a))**2.4)

def _xyz_to_lab(xyz):
    Xn, Yn, Zn = 0.95047, 1.00000, 1.08883
    x = xyz[:,0:1,:,:] / Xn
    y = xyz[:,1:2,:,:] / Yn
    z = xyz[:,2:3,:,:] / Zn
    eps = 216/24389
    kappa = 24389/27
    def f(t):
        return torch.where(t > eps, t.pow(1/3), (kappa * t + 16) / 116)
    fx, fy, fz = f(x), f(y), f(z)
    L = 116*fy - 16
    a = 500*(fx - fy)
    b = 200*(fy - fz)
    return torch.cat([L, a, b], dim=1)

def rgb_to_lab(img_bchw):
    r = img_bchw[:,0:1,:,:]; g = img_bchw[:,1:2,:,:]; b = img_bchw[:,2:3,:,:]
    rl, gl, bl = _srgb_to_linear(r), _srgb_to_linear(g), _srgb_to_linear(b)
    X = 0.4124564*rl + 0.3575761*gl + 0.1804375*bl
    Y = 0.2126729*rl + 0.7151522*gl + 0.0721750*bl
    Z = 0.0193339*rl + 0.1191920*gl + 0.9503041*bl
    xyz = torch.cat([X, Y, Z], dim=1)
    return _xyz_to_lab(xyz)

def deltaE2000(Lab1, Lab2):
    L1, a1, b1 = Lab1[:,0:1], Lab1[:,1:2], Lab1[:,2:3]
    L2, a2, b2 = Lab2[:,0:1], Lab2[:,1:2], Lab2[:,2:3]
    kL = kC = kH = 1.0

    C1 = torch.sqrt(a1*a1 + b1*b1 + 1e-12)
    C2 = torch.sqrt(a2*a2 + b2*b2 + 1e-12)
    C_ave = 0.5*(C1 + C2)
    G = 0.5*(1 - torch.sqrt((C_ave**7) / (C_ave**7 + (25.0**7))))
    a1p = (1+G)*a1
    a2p = (1+G)*a2
    C1p = torch.sqrt(a1p*a1p + b1*b1 + 1e-12)
    C2p = torch.sqrt(a2p*a2p + b2*b2 + 1e-12)

    def atan2d(y,x):
        ang = torch.rad2deg(torch.atan2(y, x))
        ang = torch.where(ang < 0, ang + 360.0, ang)
        return ang

    h1p = atan2d(b1, a1p)
    h2p = atan2d(b2, a2p)

    dLp = L2 - L1
    dCp = C2p - C1p

    dhp = h2p - h1p
    dhp = torch.where(dhp > 180.0, dhp - 360.0, dhp)
    dhp = torch.where(dhp < -180.0, dhp + 360.0, dhp)
    dHp = 2.0*torch.sqrt(C1p*C2p + 1e-12) * torch.sin(torch.deg2rad(dhp/2.0))

    Lp_ave = 0.5*(L1 + L2)
    Cp_ave = 0.5*(C1p + C2p)

    hp_sum = h1p + h2p
    hp_diff = torch.abs(h1p - h2p)
    hp_ave = torch.where((C1p*C2p).squeeze(1) == 0, hp_sum*0.0,
                         torch.where(hp_diff <= 180.0, 0.5*hp_sum,
                                     torch.where(hp_sum < 360.0, 0.5*(hp_sum + 360.0), 0.5*(hp_sum - 360.0))))
    hp_ave = hp_ave.unsqueeze(1)

    T = (1
         - 0.17*torch.cos(torch.deg2rad(hp_ave - 30.0))
         + 0.24*torch.cos(torch.deg2rad(2*hp_ave))
         + 0.32*torch.cos(torch.deg2rad(3*hp_ave + 6.0))
         - 0.20*torch.cos(torch.deg2rad(4*hp_ave - 63.0)))

    dtheta = 30.0*torch.exp(-(((hp_ave - 275.0)/25.0)**2))
    Rc = 2.0*torch.sqrt((Cp_ave**7) / (Cp_ave**7 + (25.0**7)))
    Sl = 1 + (0.015*((Lp_ave - 50.0)**2)) / torch.sqrt(20 + ((Lp_ave - 50.0)**2))
    Sc = 1 + 0.045*Cp_ave
    Sh = 1 + 0.015*Cp_ave*T
    Rt = -torch.sin(torch.deg2rad(2.0*dtheta)) * Rc

    dE = torch.sqrt(
        (dLp/(kL*Sl))**2 +
        (dCp/(kC*Sc))**2 +
        (dHp/(kH*Sh))**2 +
        Rt * (dCp/(kC*Sc)) * (dHp/(kH*Sh)) + 1e-12
    )
    return dE  # [B,1,H,W]

# ======================= Gradients & Edge IoU ===========
def scharr_grad(gray: torch.Tensor):
    kx = torch.tensor([[3, 0, -3],
                       [10, 0, -10],
                       [3, 0, -3]], dtype=gray.dtype, device=gray.device).view(1,1,3,3) / 16.0
    ky = torch.tensor([[3, 10, 3],
                       [0,  0,  0],
                       [-3,-10,-3]], dtype=gray.dtype, device=gray.device).view(1,1,3,3) / 16.0
    gx = F.conv2d(gray, kx, padding=1)
    gy = F.conv2d(gray, ky, padding=1)
    mag = torch.sqrt(gx*gx + gy*gy + 1e-12)
    return mag, gx, gy

def percentile_threshold(x: torch.Tensor, q: float):
    """
    x: [B,1,H,W]; q in [0,1] -> per-image threshold [B,1,1,1]
    Uses torch.quantile when available, else sort-index fallback.
    """
    q = float(max(0.0, min(1.0, q)))
    B = x.shape[0]
    flat = x.view(B, -1)
    try:
        th = torch.quantile(flat, q, dim=1)  # [B]
        return th.view(B,1,1,1)
    except Exception:
        N = flat.shape[1]
        k = int(q * (N - 1))
        k = 0 if k < 0 else (N - 1 if k > (N - 1) else k)
        vals, _ = torch.sort(flat, dim=1)
        th = vals[:, k]
        return th.view(B,1,1,1)

def edge_mask(gray: torch.Tensor, q: float = 0.85):
    """
    Simple Canny-like edge mask using Gaussian blur + Scharr + percentile threshold.
    gray: [B,1,H,W] in [0,1]
    returns: [B,1,H,W] binary mask in {0,1}
    """
    blur = gaussian_blur(gray, win=5, sigma=1.0)
    mag, _, _ = scharr_grad(blur)
    th = percentile_threshold(mag, q)  # [B,1,1,1]
    mask = (mag >= th).to(gray.dtype)
    return mask

def edge_iou(mask_a: torch.Tensor, mask_b: torch.Tensor):
    inter = (mask_a * mask_b)
    union = torch.clamp(mask_a + mask_b, 0, 1)
    iou = (inter.sum(dim=(2,3), keepdim=True) + 1e-6) / (union.sum(dim=(2,3), keepdim=True) + 1e-6)
    vis = torch.zeros_like(mask_a)
    vis = torch.where(inter > 0, torch.ones_like(vis), vis)
    vis = torch.where((union>0) & (inter==0), torch.full_like(vis, 0.3), vis)
    return vis.clamp(0,1), iou  # vis [B,1,H,W], scalar IoU per-image

# ======================= Wavelet (Haar) ==================
def haar_dwt2(x):
    h = torch.tensor([1.0, 1.0], device=x.device, dtype=x.dtype).view(1,1,1,2) / math.sqrt(2.0)
    g = torch.tensor([1.0,-1.0], device=x.device, dtype=x.dtype).view(1,1,1,2) / math.sqrt(2.0)
    C = x.shape[1]
    h = h.expand(C,1,1,2); g = g.expand(C,1,1,2)
    low = F.conv2d(x, h, groups=C, stride=(1,2), padding=(0,0))
    high= F.conv2d(x, g, groups=C, stride=(1,2), padding=(0,0))
    hT = h.transpose(2,3)
    gT = g.transpose(2,3)
    LL = F.conv2d(low,  hT, groups=C, stride=(2,1), padding=(0,0))
    LH = F.conv2d(low,  gT, groups=C, stride=(2,1), padding=(0,0))
    HL = F.conv2d(high, hT, groups=C, stride=(2,1), padding=(0,0))
    HH = F.conv2d(high, gT, groups=C, stride=(2,1), padding=(0,0))
    return LL, LH, HL, HH

def wavelet_band_diff(a, b, levels=3):
    baseH, baseW = a.shape[2], a.shape[3]
    accum = None
    A = a
    B = b
    for _ in range(levels):
        LA, LHa, HLa, HHa = haar_dwt2(A)
        LB, LHb, HLb, HHb = haar_dwt2(B)
        diff = (torch.abs(LHa - LHb) + torch.abs(HLa - HLb) + torch.abs(HHa - HHb)).mean(dim=1, keepdim=True)
        diff = F.interpolate(diff, size=(baseH, baseW), mode="bilinear", align_corners=False)
        accum = diff if accum is None else accum + diff
        A, B = LA, LB
    return normalize01(accum)

# ======================= Denoiser Residuum ===============
def denoiser_residual(img: torch.Tensor, win=11, sigma=2.0, gain=2.0):
    low = gaussian_blur(img, win=win, sigma=sigma)
    res = img - low
    vis = torch.clamp(0.5 + gain * res, 0.0, 1.0)
    return vis

# ======================= FFT / Spectra ===================
def fft_magnitude(x_gray: torch.Tensor, log_eps: float = 1e-3):
    x = x_gray - x_gray.mean(dim=(2,3), keepdim=True)
    fx = torch.fft.fft2(x.float())
    fx = torch.fft.fftshift(fx, dim=(-2,-1))
    mag = torch.log(torch.abs(fx) + log_eps)
    return normalize01(mag)

# ======================= Round-trip ======================
def resample(img: torch.Tensor, scale: float, kernel: str = "bicubic"):
    H, W = img.shape[2], img.shape[3]
    newH = max(1, int(round(H * scale)))
    newW = max(1, int(round(W * scale)))
    mode = "bicubic" if kernel == "bicubic" else ("bilinear" if kernel == "bilinear" else "area")
    down = F.interpolate(img, size=(newH,newW), mode=mode, align_corners=False if mode!="area" else None)
    up   = F.interpolate(down, size=(H,W), mode=mode, align_corners=False if mode!="area" else None)
    return up

# ======================= Node ============================
class MultiAnalysisNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_a": ("IMAGE",),
                "image_b": ("IMAGE",),
                "align_mode": (["resize_b_to_a", "center_crop_to_min", "pad_to_max"],),
                "color_map": (["jet", "grayscale"],),
                "invert_ssim": ("BOOLEAN", {"default": True}),
                "ms_scales": ("INT", {"default": 3, "min": 1, "max": 5}),
                "enable_lpips": ("BOOLEAN", {"default": False}),
                # Edge settings
                "edge_percentile": ("FLOAT", {"default": 0.85, "min": 0.5, "max": 0.99, "step": 0.01}),
                # Wavelet settings
                "wavelet_levels": ("INT", {"default": 3, "min": 1, "max": 4}),
                # Denoiser
                "denoise_sigma": ("FLOAT", {"default": 2.0, "min": 0.5, "max": 5.0, "step": 0.1}),
                "denoise_gain": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 10.0, "step": 0.1}),
                # Round-trip
                "roundtrip_scale": ("FLOAT", {"default": 0.5, "min": 0.25, "max": 0.95, "step": 0.05}),
                "roundtrip_kernel": (["bicubic", "bilinear", "area"],),
            }
        }

    RETURN_TYPES = (
        # Core
        "IMAGE",  # ssim_heatmap
        "IMAGE",  # dssim_heatmap
        "IMAGE",  # ms_ssim_heatmap
        "IMAGE",  # lpips_heatmap
        "IMAGE",  # diff_abs_heatmap
        "IMAGE",  # residual_visual
        "IMAGE",  # fft_orig
        "IMAGE",  # fft_upscaled
        "IMAGE",  # fft_diff
        "STRING", # ssim_text
        "STRING", # lpips_text
        # Color & ΔE2000
        "IMAGE",  # deltaE2000_heatmap
        "IMAGE",  # L_diff_heatmap
        "IMAGE",  # a_diff_heatmap
        "IMAGE",  # b_diff_heatmap
        # Gradients & edges
        "IMAGE",  # gradient_diff_heatmap
        "IMAGE",  # edge_iou_heatmap
        "STRING", # edge_iou_text
        # Wavelet
        "IMAGE",  # wavelet_band_diff_heatmap
        # Denoiser residuum
        "IMAGE",  # denoiser_residual_visual
        # Round-trip
        "IMAGE",  # roundtrip_diff_heatmap (B vs RT(B))
    )
    RETURN_NAMES = (
        "ssim_heatmap",
        "dssim_heatmap",
        "ms_ssim_heatmap",
        "lpips_heatmap",
        "diff_abs_heatmap",
        "residual_visual",
        "fft_orig",
        "fft_upscaled",
        "fft_diff",
        "ssim_text",
        "lpips_text",
        "deltaE2000_heatmap",
        "L_diff_heatmap",
        "a_diff_heatmap",
        "b_diff_heatmap",
        "gradient_diff_heatmap",
        "edge_iou_heatmap",
        "edge_iou_text",
        "wavelet_band_diff_heatmap",
        "denoiser_residual_visual",
        "roundtrip_diff_heatmap",
    )
    FUNCTION = "analyze"
    CATEGORY = NODE_CATEGORY

    def analyze(self, image_a, image_b, align_mode, color_map, invert_ssim, ms_scales, enable_lpips,
                edge_percentile, wavelet_levels, denoise_sigma, denoise_gain, roundtrip_scale, roundtrip_kernel):

        # Device/Dtype
        image_a, image_b = ensure_same_device(image_a, image_b)
        a = to_bchw(image_a)
        b = to_bchw(image_b)
        # Channels
        if a.shape[1] == 1: a = a.repeat(1,3,1,1)
        if b.shape[1] == 1: b = b.repeat(1,3,1,1)
        # Numerics
        a = a.to(dtype=torch.float32, copy=False)
        b = b.to(dtype=torch.float32, copy=False)

        # Adaptive cap defaults for SSIM-like visualizations
        SSIM_BASE_MIN = 0.02   # minimum cap (keeps identical images black)
        SSIM_QCAP     = 0.98   # percentile cap (prevents "all white" when diffs are widespread but small)

        device_type = "cuda" if a.device.type == "cuda" else "cpu"
        with torch.no_grad(), torch.autocast(device_type=device_type, enabled=False):
            # Align
            if align_mode == "resize_b_to_a":
                b = resize_like(b, a)
            elif align_mode == "center_crop_to_min":
                a, b = center_crop_to_min(a, b)
            else:
                a, b = pad_to_max(a, b)

            # ---------- SSIM & DSSIM (adaptive, zero-sensitive visualization) ----------
            ssim = ssim_map(a, b)                  # [B,1,H,W] in [0,1]
            dssim = (1.0 - ssim).clamp_(0, 1)      # difference-like
            dssim_norm = normalize_heatmap_adaptive(dssim, base_min=SSIM_BASE_MIN, qcap=SSIM_QCAP)

            # ssim_heat: if invert -> show diff (bright=change), else -> show similarity (bright=similar)
            ssim_heat = apply_colormap(dssim_norm if invert_ssim else (1.0 - dssim_norm), color_map)
            dssim_heat = apply_colormap(dssim_norm, color_map)

            # debug text additions
            l1 = torch.mean(torch.abs(a - b)).item()
            mx = torch.max(torch.abs(a - b)).item()
            ssim_text = f"SSIM mean: {float(ssim.mean().item()):.6f} | L1: {l1:.3e} | maxAbs: {mx:.3e}"

            # ---------- MS-SSIM (same adaptive logic) ----------
            ms = ms_ssim_map(a, b, scales=int(ms_scales))  # [B,1,H,W]
            ms_diff = (1.0 - ms).clamp_(0,1)
            ms_norm = normalize_heatmap_adaptive(ms_diff, base_min=SSIM_BASE_MIN, qcap=SSIM_QCAP)
            ms_heat = apply_colormap(ms_norm if invert_ssim else (1.0 - ms_norm), color_map)

            # ---------- LPIPS (optional, keeps per-image normalize) ----------
            if enable_lpips:
                lp_map, lp_text = lpips_heatmap(a, b)
                lpips_heat = apply_colormap(lp_map, color_map) if lp_map is not None else torch.zeros_like(ssim_heat)
            else:
                lpips_heat = torch.zeros_like(ssim_heat)
                lp_text = "LPIPS disabled"

            # ---------- Per-pixel abs diff ----------
            diff_abs = torch.abs(a - b).mean(dim=1, keepdim=True)
            diff_abs_heat = apply_colormap(normalize01(diff_abs), color_map)

            # ---------- Residual (signed, quick) ----------
            resid_vis = torch.clamp(0.5 + (b - a)*2.0, 0.0, 1.0)

            # ---------- FFT spectra ----------
            a_gray = (0.2989*a[:,0:1] + 0.5870*a[:,1:2] + 0.1140*a[:,2:3]).clamp_(0,1)
            b_gray = (0.2989*b[:,0:1] + 0.5870*b[:,1:2] + 0.1140*b[:,2:3]).clamp_(0,1)
            fa = fft_magnitude(a_gray, 1e-3)
            fb = fft_magnitude(b_gray, 1e-3)
            f_diff = normalize01(torch.abs(fa - fb))
            fft_a = apply_colormap(fa, color_map)
            fft_b = apply_colormap(fb, color_map)
            fft_d = apply_colormap(f_diff, color_map)

            # ---------- ΔE2000 + Lab diffs ----------
            LabA = rgb_to_lab(a.clamp(0,1))
            LabB = rgb_to_lab(b.clamp(0,1))
            dE = deltaE2000(LabA, LabB)                      # [B,1,H,W]
            dE_heat = apply_colormap(normalize01(dE), color_map)
            Ldiff = normalize01(torch.abs(LabA[:,0:1] - LabB[:,0:1]))
            adiff = normalize01(torch.abs(LabA[:,1:2] - LabB[:,1:2]))
            bdiff = normalize01(torch.abs(LabA[:,2:3] - LabB[:,2:3]))
            Ldiff_heat = apply_colormap(Ldiff, color_map)
            adiff_heat = apply_colormap(adiff, color_map)
            bdiff_heat = apply_colormap(bdiff, color_map)

            # ---------- Gradients & Edge IoU ----------
            gA, _, _ = scharr_grad(a_gray)
            gB, _, _ = scharr_grad(b_gray)
            gdiff = normalize01(torch.abs(gA - gB))
            gdiff_heat = apply_colormap(gdiff, color_map)

            mA = edge_mask(a_gray, q=float(edge_percentile))
            mB = edge_mask(b_gray, q=float(edge_percentile))
            edge_vis, iou_scalar = edge_iou(mA, mB)
            edge_heat = apply_colormap(edge_vis, color_map)
            edge_text = f"Edge IoU (percentile {edge_percentile:.2f}): {float(iou_scalar.mean().item()):.6f}"

            # ---------- Wavelet band diff ----------
            wav = wavelet_band_diff(a, b, levels=int(wavelet_levels))
            wav_heat = apply_colormap(wav, color_map)

            # ---------- Denoiser residuum ----------
            den_vis = denoiser_residual(b, win=11, sigma=float(denoise_sigma), gain=float(denoise_gain))

            # ---------- Round-trip stability ----------
            rt = resample(b, scale=float(roundtrip_scale), kernel=str(roundtrip_kernel))
            rt_diff = normalize01(torch.abs(b - rt).mean(dim=1, keepdim=True))
            rt_heat = apply_colormap(rt_diff, color_map)

        # BHWC (3ch) & return
        return (
            _ensure_bhwc3(ssim_heat),
            _ensure_bhwc3(dssim_heat),
            _ensure_bhwc3(ms_heat),
            _ensure_bhwc3(lpips_heat),
            _ensure_bhwc3(diff_abs_heat),
            _ensure_bhwc3(resid_vis),
            _ensure_bhwc3(fft_a),
            _ensure_bhwc3(fft_b),
            _ensure_bhwc3(fft_d),
            ssim_text,
            lp_text,
            _ensure_bhwc3(dE_heat),
            _ensure_bhwc3(Ldiff_heat),
            _ensure_bhwc3(adiff_heat),
            _ensure_bhwc3(bdiff_heat),
            _ensure_bhwc3(gdiff_heat),
            _ensure_bhwc3(edge_heat),
            edge_text,
            _ensure_bhwc3(wav_heat),
            _ensure_bhwc3(den_vis),
            _ensure_bhwc3(rt_heat),
        )

# ======================= Registration ====================
NODE_CLASS_MAPPINGS = {
    "MultiAnalysisNode": MultiAnalysisNode
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "MultiAnalysisNode": NODE_NAME
}