import math
import torch

class SpriteHeadStabilizeX:
    """
    Stabilize sprite animation wiggle (X only) using a Y-band (e.g. head region).

    Align frames 1..N to frame 0 by estimating horizontal shift from alpha visibility
    inside the selected Y-range.

    Methods:
      - bbox_center: leftmost/rightmost visible pixel columns -> center
      - alpha_com: alpha-weighted center-of-mass (recommended)
      - profile_corr: phase correlation on horizontal alpha profile (very robust)
      - hybrid: profile_corr with a sanity check fallback to alpha_com

    Inputs support:
      - True RGBA IMAGE tensor (C>=4) => alpha taken from channel 4
      - Or IMAGE (RGB) + MASK (ComfyUI LoadImage mask) => alpha derived from mask
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE", {}),

                # Head band
                "y_min": ("INT", {"default": 210, "min": -99999, "max": 99999, "step": 1}),
                "y_max": ("INT", {"default": 332, "min": -99999, "max": 99999, "step": 1}),

                # Alpha tolerance: visible if alpha > threshold_8bit / 255
                "alpha_threshold_8bit": ("INT", {"default": 5, "min": 0, "max": 255, "step": 1}),

                "method": (["bbox_center", "alpha_com", "profile_corr", "hybrid"], {"default": "alpha_com"}),

                # ComfyUI LoadImage produces MASK from alpha and inverts it.
                # If your mask is already alpha (0=transparent,1=opaque), set False.
                "mask_is_inverted": ("BOOLEAN", {"default": True}),

                # Optional safety clamps/smoothing
                "max_abs_shift": ("INT", {"default": 0, "min": 0, "max": 99999, "step": 1}),
                "temporal_median": ("INT", {"default": 1, "min": 1, "max": 99, "step": 1}),

                # Hybrid sanity check: if corr shift differs from COM shift by more than this,
                # use COM shift instead.
                "hybrid_tolerance_px": ("INT", {"default": 8, "min": 0, "max": 99999, "step": 1}),
            },
            "optional": {
                "mask": ("MASK", {}),
            }
        }

    RETURN_TYPES = ("IMAGE", "MASK", "STRING")
    RETURN_NAMES = ("images", "mask", "shifts_x")
    FUNCTION = "stabilize"
    CATEGORY = "image/sprite"
    SEARCH_ALIASES = ["wiggle stabilize", "sprite stabilize", "head stabilize", "animation stabilize", "sprite jitter fix"]

    # ---------- helpers ----------

    def _get_alpha(self, images: torch.Tensor, mask: torch.Tensor | None, mask_is_inverted: bool) -> torch.Tensor:
        """
        Returns alpha in [0..1], shape [B,H,W].
        """
        if images.dim() != 4:
            raise ValueError(f"images must have shape [B,H,W,C], got {tuple(images.shape)}")
        B, H, W, C = images.shape

        if C >= 4:
            return images[..., 3]

        if mask is None:
            raise ValueError("Need RGBA images (C>=4) OR provide a MASK input.")

        if mask.dim() == 2:
            mask = mask.unsqueeze(0)
        if mask.dim() != 3:
            raise ValueError(f"mask must have shape [B,H,W] or [H,W], got {tuple(mask.shape)}")

        if mask.shape[1] != H or mask.shape[2] != W:
            raise ValueError(f"mask H/W must match images; mask={tuple(mask.shape)} images={tuple(images.shape)}")

        if mask.shape[0] == 1 and B > 1:
            mask = mask.repeat(B, 1, 1)
        if mask.shape[0] != B:
            raise ValueError(f"mask batch must match images batch; mask B={mask.shape[0]} images B={B}")

        alpha = 1.0 - mask if mask_is_inverted else mask
        return alpha

    def _clamp_y(self, H: int, y_min: int, y_max: int) -> tuple[int, int]:
        y0 = int(y_min)
        y1 = int(y_max)
        if y1 < y0:
            y0, y1 = y1, y0
        y0 = max(0, min(H - 1, y0))
        y1 = max(0, min(H - 1, y1))
        return y0, y1

    def _bbox_center_x(self, alpha_hw: torch.Tensor, thr: float) -> float | None:
        """
        alpha_hw: [H,W]
        Returns center X using leftmost/rightmost visible columns, or None if empty.
        """
        # visible: [H,W]
        visible = alpha_hw > thr
        cols = visible.any(dim=0)  # [W]
        if not bool(cols.any()):
            return None
        W = alpha_hw.shape[1]
        left = int(torch.argmax(cols.float()).item())
        right = int((W - 1) - torch.argmax(torch.flip(cols, dims=[0]).float()).item())
        return (left + right) / 2.0

    def _com_center_x(self, alpha_hw: torch.Tensor, thr: float) -> float | None:
        """
        alpha_hw: [H,W]
        Alpha-weighted center-of-mass of X within visible area, or None if empty.
        """
        W = alpha_hw.shape[1]
        weights = alpha_hw
        if thr > 0:
            weights = weights * (weights > thr)

        profile = weights.sum(dim=0)  # [W]
        total = float(profile.sum().item())
        if total <= 0.0:
            return None

        x = torch.arange(W, device=alpha_hw.device, dtype=profile.dtype)
        center = float((profile * x).sum().item() / total)
        return center

    def _phase_corr_shift_x(self, alpha_hw: torch.Tensor, ref_profile: torch.Tensor, thr: float) -> int | None:
        """
        Estimate integer shift to APPLY to current frame (X) so it matches reference.
        Uses 1D phase correlation on horizontal alpha profile.
        Returns shift_x (int), or None if empty.
        """
        weights = alpha_hw
        if thr > 0:
            weights = weights * (weights > thr)

        prof = weights.sum(dim=0).float()
        if float(prof.sum().item()) <= 0.0:
            return None

        # Remove DC component
        prof = prof - prof.mean()
        ref = ref_profile

        # Phase correlation
        F = torch.fft.rfft(prof)
        R = torch.fft.rfft(ref)
        cps = F * torch.conj(R)
        cps = cps / (torch.abs(cps) + 1e-9)
        corr = torch.fft.irfft(cps, n=prof.numel())
        peak = int(torch.argmax(corr).item())

        W = prof.numel()
        lag = peak if peak <= W // 2 else peak - W  # lag = "current is shifted by lag relative to ref"
        shift_x = -lag  # apply negative to align to ref
        return int(shift_x)

    def _shift_frame_x(self, img_hwc: torch.Tensor, shift_x: int) -> torch.Tensor:
        """
        img_hwc: [H,W,C]
        shift_x: int (positive -> move right)
        """
        H, W, C = img_hwc.shape
        out = torch.zeros_like(img_hwc)
        if shift_x == 0:
            return img_hwc
        if abs(shift_x) >= W:
            return out

        if shift_x > 0:
            out[:, shift_x:, :] = img_hwc[:, : W - shift_x, :]
        else:
            sx = -shift_x
            out[:, : W - sx, :] = img_hwc[:, sx:, :]
        return out

    def _shift_mask_x(self, m_hw: torch.Tensor, shift_x: int, fill_val: float) -> torch.Tensor:
        """
        m_hw: [H,W]
        """
        H, W = m_hw.shape
        out = torch.full_like(m_hw, fill_val)
        if shift_x == 0:
            return m_hw
        if abs(shift_x) >= W:
            return out
        if shift_x > 0:
            out[:, shift_x:] = m_hw[:, : W - shift_x]
        else:
            sx = -shift_x
            out[:, : W - sx] = m_hw[:, sx:]
        return out

    def _median_smooth(self, shifts: list[int], window: int) -> list[int]:
        """
        Median filter over shifts with odd window size. Keeps shifts[0] unchanged.
        """
        if window <= 1 or len(shifts) <= 2:
            return shifts
        w = int(window)
        if w % 2 == 0:
            w += 1
        r = w // 2
        out = shifts[:]
        out[0] = shifts[0]
        n = len(shifts)
        for i in range(1, n):
            lo = max(1, i - r)
            hi = min(n, i + r + 1)
            vals = sorted(shifts[lo:hi])
            out[i] = vals[len(vals) // 2]
        return out

    # ---------- main ----------

    def stabilize(
        self,
        images: torch.Tensor,
        y_min: int = 210,
        y_max: int = 332,
        alpha_threshold_8bit: int = 5,
        method: str = "alpha_com",
        mask_is_inverted: bool = True,
        max_abs_shift: int = 0,
        temporal_median: int = 1,
        hybrid_tolerance_px: int = 8,
        mask: torch.Tensor | None = None,
    ):
        if not torch.is_tensor(images):
            raise TypeError("images must be a torch.Tensor")
        if images.dim() != 4:
            raise ValueError(f"images must have shape [B,H,W,C], got {tuple(images.shape)}")

        B, H, W, C = images.shape
        if B < 1:
            raise ValueError("images batch is empty")

        alpha = self._get_alpha(images, mask, mask_is_inverted)  # [B,H,W]
        y0, y1 = self._clamp_y(H, y_min, y_max)
        thr = float(alpha_threshold_8bit) / 255.0

        roi_alpha = alpha[:, y0:y1 + 1, :]  # [B, Hr, W]

        # Reference (frame 0)
        ref_roi = roi_alpha[0]  # [Hr,W]

        # Prepare reference for methods
        ref_center_bbox = None
        ref_center_com = None
        ref_profile = None

        if method in ("bbox_center", "hybrid"):
            ref_center_bbox = self._bbox_center_x(ref_roi, thr)
        if method in ("alpha_com", "hybrid"):
            ref_center_com = self._com_center_x(ref_roi, thr)
        if method in ("profile_corr", "hybrid"):
            # reference profile for phase correlation
            w = ref_roi
            if thr > 0:
                w = w * (w > thr)
            ref_profile = w.sum(dim=0).float()
            ref_profile = ref_profile - ref_profile.mean()

        # Fallback reference center if missing
        if ref_center_bbox is None and ref_center_com is None and ref_profile is None:
            # Nothing visible even in reference head region; do nothing.
            out_mask = None
            if mask is not None:
                out_mask = mask if mask.dim() == 3 else mask.unsqueeze(0)
            elif C >= 4:
                a = images[..., 3]
                out_mask = (1.0 - a) if mask_is_inverted else a
            else:
                fill_val = 1.0 if mask_is_inverted else 0.0
                out_mask = torch.full((B, H, W), fill_val, device=images.device, dtype=images.dtype)

            return (images, out_mask, "[0]" if B == 1 else str([0] * B))

        # For center-based methods, pick a reference center
        # Preference: COM, else BBOX, else image center
        if ref_center_com is not None:
            ref_center = ref_center_com
        elif ref_center_bbox is not None:
            ref_center = ref_center_bbox
        else:
            ref_center = W / 2.0

        shifts = [0] * B
        shifts[0] = 0  # frame 0 stays

        for i in range(1, B):
            a_hw = roi_alpha[i]

            shift_i = 0

            if method == "bbox_center":
                c = self._bbox_center_x(a_hw, thr)
                if c is None:
                    shift_i = 0
                else:
                    shift_i = int(round(ref_center - c))

            elif method == "alpha_com":
                c = self._com_center_x(a_hw, thr)
                if c is None:
                    shift_i = 0
                else:
                    shift_i = int(round(ref_center - c))

            elif method == "profile_corr":
                s = self._phase_corr_shift_x(a_hw, ref_profile, thr)  # already int shift to APPLY
                shift_i = 0 if s is None else int(s)

            elif method == "hybrid":
                # corr shift
                s_corr = self._phase_corr_shift_x(a_hw, ref_profile, thr) if ref_profile is not None else None

                # com shift
                c = self._com_center_x(a_hw, thr)
                s_com = None if c is None else int(round(ref_center - c))

                if s_corr is None and s_com is None:
                    shift_i = 0
                elif s_corr is None:
                    shift_i = int(s_com)
                elif s_com is None:
                    shift_i = int(s_corr)
                else:
                    if abs(int(s_corr) - int(s_com)) > int(hybrid_tolerance_px):
                        shift_i = int(s_com)
                    else:
                        shift_i = int(s_corr)

            else:
                raise ValueError(f"Unknown method: {method}")

            # Clamp extreme shifts if requested
            if max_abs_shift and max_abs_shift > 0:
                shift_i = int(max(-max_abs_shift, min(max_abs_shift, shift_i)))

            shifts[i] = shift_i

        # Optional temporal median smoothing (keeps shifts[0] anchored)
        shifts = self._median_smooth(shifts, int(temporal_median))

        # Apply per-frame shifts
        out_images = torch.zeros_like(images)

        # Output mask handling:
        # - If input mask provided: shift it
        # - Else if RGBA: derive from shifted alpha
        # - Else: produce blank
        out_mask = None
        in_mask_bhw = None
        if mask is not None:
            in_mask_bhw = mask
            if in_mask_bhw.dim() == 2:
                in_mask_bhw = in_mask_bhw.unsqueeze(0)
            if in_mask_bhw.shape[0] == 1 and B > 1:
                in_mask_bhw = in_mask_bhw.repeat(B, 1, 1)

            fill_val = 1.0 if mask_is_inverted else 0.0
            out_mask = torch.full_like(in_mask_bhw, fill_val)

        for i in range(B):
            sx = int(shifts[i])
            out_images[i] = self._shift_frame_x(images[i], sx)

            if out_mask is not None and in_mask_bhw is not None:
                fill_val = 1.0 if mask_is_inverted else 0.0
                out_mask[i] = self._shift_mask_x(in_mask_bhw[i], sx, fill_val)

        if out_mask is None:
            if out_images.shape[-1] >= 4:
                a = out_images[..., 3]
                out_mask = (1.0 - a) if mask_is_inverted else a
            else:
                fill_val = 1.0 if mask_is_inverted else 0.0
                out_mask = torch.full((B, H, W), fill_val, device=images.device, dtype=images.dtype)

        shifts_str = str(shifts)
        return (out_images, out_mask, shifts_str)


NODE_CLASS_MAPPINGS = {
    "SpriteHeadStabilizeX": SpriteHeadStabilizeX,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SpriteHeadStabilizeX": "Sprite Head Stabilize X (Batch)",
}