1#!/usr/bin/env python3
2"""
3GIF Builder - Core module for assembling frames into GIFs optimized for Slack.
4 
5This module provides the main interface for creating GIFs from programmatically
6generated frames, with automatic optimization for Slack's requirements.
7"""
8 
9from pathlib import Path
10from typing import Optional
11 
12import imageio.v3 as imageio
13import numpy as np
14from PIL import Image
15 
16 
17class GIFBuilder:
18    """Builder for creating optimized GIFs from frames."""
19 
20    def __init__(self, width: int = 480, height: int = 480, fps: int = 15):
21        """
22        Initialize GIF builder.
23 
24        Args:
25            width: Frame width in pixels
26            height: Frame height in pixels
27            fps: Frames per second
28        """
29        self.width = width
30        self.height = height
31        self.fps = fps
32        self.frames: list[np.ndarray] = []
33 
34    def add_frame(self, frame: np.ndarray | Image.Image):
35        """
36        Add a frame to the GIF.
37 
38        Args:
39            frame: Frame as numpy array or PIL Image (will be converted to RGB)
40        """
41        if isinstance(frame, Image.Image):
42            frame = np.array(frame.convert("RGB"))
43 
44        # Ensure frame is correct size
45        if frame.shape[:2] != (self.height, self.width):
46            pil_frame = Image.fromarray(frame)
47            pil_frame = pil_frame.resize(
48                (self.width, self.height), Image.Resampling.LANCZOS
49            )
50            frame = np.array(pil_frame)
51 
52        self.frames.append(frame)
53 
54    def add_frames(self, frames: list[np.ndarray | Image.Image]):
55        """Add multiple frames at once."""
56        for frame in frames:
57            self.add_frame(frame)
58 
59    def optimize_colors(
60        self, num_colors: int = 128, use_global_palette: bool = True
61    ) -> list[np.ndarray]:
62        """
63        Reduce colors in all frames using quantization.
64 
65        Args:
66            num_colors: Target number of colors (8-256)
67            use_global_palette: Use a single palette for all frames (better compression)
68 
69        Returns:
70            List of color-optimized frames
71        """
72        optimized = []
73 
74        if use_global_palette and len(self.frames) > 1:
75            # Create a global palette from all frames
76            # Sample frames to build palette
77            sample_size = min(5, len(self.frames))
78            sample_indices = [
79                int(i * len(self.frames) / sample_size) for i in range(sample_size)
80            ]
81            sample_frames = [self.frames[i] for i in sample_indices]
82 
83            # Combine sample frames into a single image for palette generation
84            # Flatten each frame to get all pixels, then stack them
85            all_pixels = np.vstack(
86                [f.reshape(-1, 3) for f in sample_frames]
87            )  # (total_pixels, 3)
88 
89            # Create a properly-shaped RGB image from the pixel data
90            # We'll make a roughly square image from all the pixels
91            total_pixels = len(all_pixels)
92            width = min(512, int(np.sqrt(total_pixels)))  # Reasonable width, max 512
93            height = (total_pixels + width - 1) // width  # Ceiling division
94 
95            # Pad if necessary to fill the rectangle
96            pixels_needed = width * height
97            if pixels_needed > total_pixels:
98                padding = np.zeros((pixels_needed - total_pixels, 3), dtype=np.uint8)
99                all_pixels = np.vstack([all_pixels, padding])
100 
101            # Reshape to proper RGB image format (H, W, 3)
102            img_array = (
103                all_pixels[:pixels_needed].reshape(height, width, 3).astype(np.uint8)
104            )
105            combined_img = Image.fromarray(img_array, mode="RGB")
106 
107            # Generate global palette
108            global_palette = combined_img.quantize(colors=num_colors, method=2)
109 
110            # Apply global palette to all frames
111            for frame in self.frames:
112                pil_frame = Image.fromarray(frame)
113                quantized = pil_frame.quantize(palette=global_palette, dither=1)
114                optimized.append(np.array(quantized.convert("RGB")))
115        else:
116            # Use per-frame quantization
117            for frame in self.frames:
118                pil_frame = Image.fromarray(frame)
119                quantized = pil_frame.quantize(colors=num_colors, method=2, dither=1)
120                optimized.append(np.array(quantized.convert("RGB")))
121 
122        return optimized
123 
124    def deduplicate_frames(self, threshold: float = 0.9995) -> int:
125        """
126        Remove duplicate or near-duplicate consecutive frames.
127 
128        Args:
129            threshold: Similarity threshold (0.0-1.0). Higher = more strict (0.9995 = nearly identical).
130                      Use 0.9995+ to preserve subtle animations, 0.98 for aggressive removal.
131 
132        Returns:
133            Number of frames removed
134        """
135        if len(self.frames) < 2:
136            return 0
137 
138        deduplicated = [self.frames[0]]
139        removed_count = 0
140 
141        for i in range(1, len(self.frames)):
142            # Compare with previous frame
143            prev_frame = np.array(deduplicated[-1], dtype=np.float32)
144            curr_frame = np.array(self.frames[i], dtype=np.float32)
145 
146            # Calculate similarity (normalized)
147            diff = np.abs(prev_frame - curr_frame)
148            similarity = 1.0 - (np.mean(diff) / 255.0)
149 
150            # Keep frame if sufficiently different
151            # High threshold (0.9995+) means only remove nearly identical frames
152            if similarity < threshold:
153                deduplicated.append(self.frames[i])
154            else:
155                removed_count += 1
156 
157        self.frames = deduplicated
158        return removed_count
159 
160    def save(
161        self,
162        output_path: str | Path,
163        num_colors: int = 128,
164        optimize_for_emoji: bool = False,
165        remove_duplicates: bool = False,
166    ) -> dict:
167        """
168        Save frames as optimized GIF for Slack.
169 
170        Args:
171            output_path: Where to save the GIF
172            num_colors: Number of colors to use (fewer = smaller file)
173            optimize_for_emoji: If True, optimize for emoji size (128x128, fewer colors)
174            remove_duplicates: If True, remove duplicate consecutive frames (opt-in)
175 
176        Returns:
177            Dictionary with file info (path, size, dimensions, frame_count)
178        """
179        if not self.frames:
180            raise ValueError("No frames to save. Add frames with add_frame() first.")
181 
182        output_path = Path(output_path)
183 
184        # Remove duplicate frames to reduce file size
185        if remove_duplicates:
186            removed = self.deduplicate_frames(threshold=0.9995)
187            if removed > 0:
188                print(
189                    f"  Removed {removed} nearly identical frames (preserved subtle animations)"
190                )
191 
192        # Optimize for emoji if requested
193        if optimize_for_emoji:
194            if self.width > 128 or self.height > 128:
195                print(
196                    f"  Resizing from {self.width}x{self.height} to 128x128 for emoji"
197                )
198                self.width = 128
199                self.height = 128
200                # Resize all frames
201                resized_frames = []
202                for frame in self.frames:
203                    pil_frame = Image.fromarray(frame)
204                    pil_frame = pil_frame.resize((128, 128), Image.Resampling.LANCZOS)
205                    resized_frames.append(np.array(pil_frame))
206                self.frames = resized_frames
207            num_colors = min(num_colors, 48)  # More aggressive color limit for emoji
208 
209            # More aggressive FPS reduction for emoji
210            if len(self.frames) > 12:
211                print(
212                    f"  Reducing frames from {len(self.frames)} to ~12 for emoji size"
213                )
214                # Keep every nth frame to get close to 12 frames
215                keep_every = max(1, len(self.frames) // 12)
216                self.frames = [
217                    self.frames[i] for i in range(0, len(self.frames), keep_every)
218                ]
219 
220        # Optimize colors with global palette
221        optimized_frames = self.optimize_colors(num_colors, use_global_palette=True)
222 
223        # Calculate frame duration in milliseconds
224        frame_duration = 1000 / self.fps
225 
226        # Save GIF
227        imageio.imwrite(
228            output_path,
229            optimized_frames,
230            duration=frame_duration,
231            loop=0,  # Infinite loop
232        )
233 
234        # Get file info
235        file_size_kb = output_path.stat().st_size / 1024
236        file_size_mb = file_size_kb / 1024
237 
238        info = {
239            "path": str(output_path),
240            "size_kb": file_size_kb,
241            "size_mb": file_size_mb,
242            "dimensions": f"{self.width}x{self.height}",
243            "frame_count": len(optimized_frames),
244            "fps": self.fps,
245            "duration_seconds": len(optimized_frames) / self.fps,
246            "colors": num_colors,
247        }
248 
249        # Print info
250        print(f"\n✓ GIF created successfully!")
251        print(f"  Path: {output_path}")
252        print(f"  Size: {file_size_kb:.1f} KB ({file_size_mb:.2f} MB)")
253        print(f"  Dimensions: {self.width}x{self.height}")
254        print(f"  Frames: {len(optimized_frames)} @ {self.fps} fps")
255        print(f"  Duration: {info['duration_seconds']:.1f}s")
256        print(f"  Colors: {num_colors}")
257 
258        # Size info
259        if optimize_for_emoji:
260            print(f"  Optimized for emoji (128x128, reduced colors)")
261        if file_size_mb > 1.0:
262            print(f"\n  Note: Large file size ({file_size_kb:.1f} KB)")
263            print("  Consider: fewer frames, smaller dimensions, or fewer colors")
264 
265        return info
266 
267    def clear(self):
268        """Clear all frames (useful for creating multiple GIFs)."""
269        self.frames = []
270