# File: ComfyUI/custom_nodes/batch_camera_shake.py

import math
import random
import numpy as np
import torch
import torch.nn.functional as F

class BatchCameraShake:
    """
    Apply subtle per-frame affine transforms (translate/rotate/scale) to an IMAGE batch.
    Designed for 'handheld' micro-movement.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE", {
                    "tooltip": "Input image batch to apply camera shake effect to"
                }),
                "seed": ("INT", {
                    "default": 0, "min": 0, "max": 2**31-1,
                    "tooltip": "Random seed for reproducible shake patterns. Same seed = same shake motion"
                }),
                "translate_px": ("FLOAT", {
                    "default": 2.0, "min": 0.0, "max": 50.0, "step": 0.1,
                    "tooltip": "Maximum horizontal/vertical movement in pixels. Higher = more pronounced shake"
                }),
                "rotate_deg": ("FLOAT", {
                    "default": 0.15, "min": 0.0, "max": 10.0, "step": 0.01,
                    "tooltip": "Maximum rotation angle in degrees. Subtle values (0.1-0.3) look natural"
                }),
                "scale_pct": ("FLOAT", {
                    "default": 0.5, "min": 0.0, "max": 10.0, "step": 0.1,
                    "tooltip": "Maximum zoom variation as a percentage. Simulates camera moving closer/further"
                }),
                "freq": ("FLOAT", {
                    "default": 1.0, "min": 0.1, "max": 5.0, "step": 0.1,
                    "tooltip": "Shake frequency multiplier. Higher = faster, more frequent shake changes"
                }),
                "smoothing": ("FLOAT", {
                    "default": 0.85, "min": 0.0, "max": 0.99, "step": 0.01,
                    "tooltip": "Motion smoothness. Higher = smoother, floaty movement. Lower = jittery, nervous shake"
                }),
                "distribution": (["even", "center", "ramp_up", "ramp_down", "ease_in", "ease_out", "ease_in_out"], {
                    "default": "even",
                    "tooltip": "Intensity distribution over time. even=uniform, center=peaks in middle, ease_out=settles to origin"
                }),
                "min_intensity": ("FLOAT", {
                    "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01,
                    "tooltip": "Minimum intensity as fraction of maximum (0-1). Only active when distribution is not 'even'. 0=full range, 1=no variation"
                }),
                "border": (["reflect", "replicate", "zeros"], {
                    "default": "reflect",
                    "tooltip": "Edge handling: reflect mirrors edges, replicate extends edge pixels, zeros fills with black"
                }),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "shake"
    CATEGORY = "image/transform"

    def _smooth_noise(self, n, rng, smoothing):
        # Simple 1D noise smoothed by exponential moving average
        x = np.array([rng.uniform(-1.0, 1.0) for _ in range(n)], dtype=np.float32)
        y = np.zeros_like(x)
        a = smoothing
        y[0] = x[0]
        for i in range(1, n):
            y[i] = a * y[i-1] + (1 - a) * x[i]
        # Normalize using standard deviation for consistent amplitude across seeds
        # This ensures different seeds produce similar shake intensity
        std = float(np.std(y)) if n > 1 else 1.0
        if std < 1e-6:
            std = 1.0
        # Scale so ~95% of values fall within [-1, 1] (2 standard deviations)
        return np.clip(y / (2.0 * std), -1.0, 1.0)

    def _distribution_envelope(self, n, mode, min_envelope=0.5):
        """
        Create an intensity envelope based on the distribution mode.
        Returns array of shape (n,) with values in [min_envelope, 1.0].
        The min_envelope parameter controls the lowest point as a fraction of maximum.
        When mode is 'even', min_envelope is ignored and all values are 1.0.
        """
        if n <= 1:
            return np.ones(n, dtype=np.float32)
        
        # Normalized position from 0.0 (first frame) to 1.0 (last frame)
        t = np.linspace(0.0, 1.0, n, dtype=np.float32)
        
        if mode == "even":
            # Even distribution ignores min_envelope - always uniform at 1.0
            return np.ones(n, dtype=np.float32)
        elif mode == "center":
            # Bell curve: peaks at middle
            envelope = np.sin(np.pi * t)
        elif mode == "ramp_up":
            # Linear increase
            envelope = t
        elif mode == "ramp_down":
            # Linear decrease
            envelope = 1.0 - t
        elif mode == "ease_in":
            # Quadratic: slow start, accelerates
            envelope = t ** 2
        elif mode == "ease_out":
            # Quadratic: fast start, decelerates (settles)
            envelope = (1.0 - t) ** 2
        elif mode == "ease_in_out":
            # Smooth S-curve
            envelope = 0.5 * (1.0 - np.cos(np.pi * t))
        else:
            # Fallback to even
            return np.ones(n, dtype=np.float32)
        
        # Scale from [0, 1] to [min_envelope, 1] so minimum is min_envelope of maximum
        return min_envelope + (1.0 - min_envelope) * envelope

    def shake(self, images, seed, translate_px, rotate_deg, scale_pct, freq, smoothing, distribution, min_intensity, border):
        """
        images: torch float32 in [0,1], shape [B,H,W,C] (ComfyUI IMAGE)
        freq: higher values = faster/more frequent shake changes
        distribution: intensity curve over time (even, center, ramp_up, etc.)
        min_intensity: minimum intensity as fraction of max (ignored when distribution='even')
        """
        device = images.device
        B, H, W, C = images.shape

        rng = random.Random(int(seed))

        # Generate smooth per-frame signals for tx, ty, rot, scale
        # freq controls how quickly shake changes: higher freq = more variation over time
        # We generate noise at an effective length, then interpolate to B frames
        effective_len = max(int(B * freq), 2)
        
        nx = self._smooth_noise(effective_len, rng, smoothing)
        ny = self._smooth_noise(effective_len, rng, smoothing)
        nr = self._smooth_noise(effective_len, rng, smoothing)
        ns = self._smooth_noise(effective_len, rng, smoothing)

        # Convert user amplitudes into actual values
        tx_raw = (translate_px * nx).astype(np.float32)
        ty_raw = (translate_px * ny).astype(np.float32)
        rot_raw = (rotate_deg * nr).astype(np.float32)
        sc_raw = ((scale_pct / 100.0) * ns).astype(np.float32)  # Store as offset from 1.0

        # Interpolate from effective_len samples to B output frames
        # This ensures smooth transitions regardless of freq value
        u = np.linspace(0, effective_len - 1, B, dtype=np.float32)
        
        def lerp(a):
            i0 = np.floor(u).astype(np.int32)
            i1 = np.clip(i0 + 1, 0, len(a) - 1)
            f = (u - i0).astype(np.float32)
            return a[i0] * (1 - f) + a[i1] * f
        
        tx, ty, rot, sc_offset = lerp(tx_raw), lerp(ty_raw), lerp(rot_raw), lerp(sc_raw)
        
        # Apply distribution envelope to shape intensity over time
        envelope = self._distribution_envelope(B, distribution, min_intensity)
        tx = tx * envelope
        ty = ty * envelope
        rot = rot * envelope
        sc = 1.0 + (sc_offset * envelope)  # Scale returns to 1.0 where envelope is 0

        # Prepare for grid_sample: NCHW
        x = images.permute(0, 3, 1, 2).contiguous()

        # Border handling
        if border in ("reflect", "replicate"):
            pad = int(math.ceil(max(translate_px, 2.0) + max(W, H) * (scale_pct / 100.0)))
            # keep pad reasonable
            pad = max(4, min(pad, 256))
            mode = "reflect" if border == "reflect" else "replicate"
            x = F.pad(x, (pad, pad, pad, pad), mode=mode)
            Hp, Wp = x.shape[-2], x.shape[-1]
        else:
            pad = 0
            Hp, Wp = H, W

        # Build per-frame affine matrices in normalized coords
        # grid_sample expects theta mapping output -> input in normalized [-1,1] coords.
        thetas = []
        for i in range(B):
            a = math.radians(float(rot[i]))
            s = float(sc[i])

            cos_a = math.cos(a) / s
            sin_a = math.sin(a) / s

            # translation in pixels -> normalized
            # Note: after padding, use Wp/Hp
            tx_n = 2.0 * float(tx[i]) / max(Wp - 1, 1)
            ty_n = 2.0 * float(ty[i]) / max(Hp - 1, 1)

            theta = torch.tensor([
                [cos_a, -sin_a, tx_n],
                [sin_a,  cos_a, ty_n],
            ], dtype=torch.float32, device=device)
            thetas.append(theta)

        theta = torch.stack(thetas, dim=0)  # [B,2,3]

        # Create sampling grid and apply transform
        grid = F.affine_grid(theta, size=x.shape, align_corners=False)
        y = F.grid_sample(
            x, grid,
            mode="bilinear",
            padding_mode="zeros",  # already padded if needed
            align_corners=False
        )

        # Remove padding and return to ComfyUI IMAGE format
        if pad > 0:
            y = y[:, :, pad:pad+H, pad:pad+W]

        out = y.permute(0, 2, 3, 1).contiguous().clamp(0.0, 1.0)
        return (out,)

NODE_CLASS_MAPPINGS = {
    "BatchCameraShake": BatchCameraShake
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BatchCameraShake": "Batch Camera Shake (Subtle)"
}