utils.py

import shutil
import datetime
import yaml
import sys
from pathlib import Path
import jax
import jax.numpy as jnp
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import seaborn as sns
from PIL import Image, ImageFont, ImageDraw

sns.set_theme(style="white", context="talk")
plt.rcParams['savefig.facecolor'] = 'white'

def count_trainable_params(model):
    """Counts the number of trainable float parameters in an Equinox module."""
    def count_params(x):
        if isinstance(x, jnp.ndarray) and x.dtype in [jnp.float32, jnp.float64]:
            return x.size
        return 0
    param_counts = jax.tree_util.tree_map(count_params, model)
    return sum(jax.tree_util.tree_leaves(param_counts))

def setup_run_dir(phase_name, config, train=True, base_dir="runs"):
    """
    Sets up the directory for the current phase.
    If train=True, creates a timestamped folder, copies the calling script,
    dumps the config to yaml, and returns the path.
    """

    ## Seriously warn the user that phase_1 should be run from the root project directory, while phase_2 and 3 from the runs/xx, directory. Use alrm emojis
    if phase_name == "phase_1":
        print("⚠️⚠️⚠️ WARNING: We recommend runing phase_1 from the root project directory ⚠️⚠️⚠️", flush=True)
    else:
        print(f"⚠️⚠️⚠️ WARNING: We recommend running {phase_name} from the run directory created by phase_1 ⚠️⚠️⚠️", flush=True)

    ## If phase 2 or 3, do nothing, return ./
    if not train or phase_name in ["phase_2", "phase_3", "phase_4"]:
    # if not train:
        # data_dir = Path(config["data_path"])
        # config["data_path"] = "../../" + str(data_dir.name)

        run_path = Path("./")

    # if train:
    # if train and phase_name not in ["phase_2", "phase_3"]:
    else:
        timestamp = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
        run_path = Path(base_dir) / timestamp
        run_path.mkdir(parents=True, exist_ok=True)

        (run_path / "artefacts").mkdir(exist_ok=True)
        (run_path / "plots").mkdir(exist_ok=True)
        
        with open(run_path / "config.yaml", 'w') as f:
            # yaml.dump(config, f, default_flow_style=False)

            ## While dumping the config, set the data_path to ../../old_data_path
            config_to_dump = config.copy()
            if "data_path" in config_to_dump:
                data_dir = Path(config_to_dump["data_path"])
                config_to_dump["data_path"] = "../../" + str(data_dir.name)
            yaml.dump(config_to_dump, f, default_flow_style=False)

        # 1. Handle current_script but ignore ipykernel_launcher
        current_script = Path(sys.argv[0])
        files_to_copy = [
            "utils.py", "loaders.py", "models.py", "phase1.py",
            "phase2.py", "phase3.py"
        ]

        # if current_script.exists() and current_script.is_file() and "ipykernel_launcher" not in current_script.name:
        #     files_to_copy.append(current_script.name)

        # 2. Copy the files and use a set() to avoid trying to copy the same file twice
        for fname in set(files_to_copy):
            src_file = Path(fname)
            if src_file.exists() and src_file.is_file():
                shutil.copy(src_file, run_path / src_file.name)
            elif fname in ["phase_1.py", "phase1.py"]:
                # Optional: Print a warning so you know exactly why it's failing if it still doesn't copy
                pass 

    return run_path

def get_coords_grid(H, W):
    """Generates a normalised coordinate grid for the INR."""
    y_coords = jnp.linspace(-1, 1, H)
    x_coords = jnp.linspace(-1, 1, W)
    X_grid, Y_grid = jnp.meshgrid(x_coords, y_coords)
    return jnp.stack([X_grid, Y_grid], axis=-1)



def plot_videos(video, ref_video=None, plot_ref=True, show_titles=True, show_labels=True, forecast_start=None, 
                vmin=None, vmax=None, save_name=None, 
                wspace=0.05, hspace=0.02, forecast_gap=0.2, 
                save_video=False, video_gap=5, show_borders=False, corner_radius=5,
                no_rescale=True, cmap='coolwarm', row_height="auto", gif_scale=4):
    """
    Plots a camera-ready rollout of ground truth and predicted video frames, 
    and saves a high-res, properly scaled GIF.
    """
    with plt.rc_context({
        'font.family': 'sans-serif',
        'font.sans-serif': ['Helvetica Neue', 'Helvetica', 'Arial', 'DejaVu Sans'],
        'font.size': 18,
        'pdf.fonttype': 42,
        'ps.fonttype': 42
    }):
        
        nb_frames = video.shape[0]
        C = video.shape[-1]
        
        if plot_ref and ref_video is None:
            raise ValueError("ref_video must be provided if plot_ref is True.")

        rescale = False
        if plot_ref and ref_video[..., :C].min() < -0.5:
            rescale = True
            ref_video = (ref_video + 1.0) / 2.0
        elif not plot_ref and video.min() < -0.5:
            rescale = True
        
        if no_rescale:
            rescale = False

        nrows = 2 if plot_ref else 1
        has_gap = forecast_start is not None and 1 < forecast_start <= nb_frames
        ncols = nb_frames + 1 if has_gap else nb_frames
        
        width_ratios = [1.0] * ncols
        spacer_col = -1
        if has_gap:
            spacer_col = forecast_start - 1
            width_ratios[spacer_col] = forecast_gap

        # Base width calculation
        fig_width = (nb_frames + (forecast_gap if has_gap else 0.0)) * 1.5
        
        if row_height == "auto":
            H, W = video.shape[1:3]
            aspect = H / W
            calculated_row_height = aspect * 1.2 if show_titles else aspect * 1.5
            title_buffer = 0.5 if show_titles else 0.1
            fig_height = (nrows * calculated_row_height) + title_buffer
        else:
            fig_height = nrows * float(row_height)
        
        fig = plt.figure(figsize=(fig_width, fig_height))
        gs = fig.add_gridspec(nrows, ncols, wspace=wspace, hspace=hspace, width_ratios=width_ratios)
        axes = np.empty((nrows, ncols), dtype=object)
        for r in range(nrows):
            for c in range(ncols):
                axes[r, c] = fig.add_subplot(gs[r, c])

        # Establish Global Min/Max
        if vmin is None or vmax is None:
            if plot_ref:
                global_min = ref_video.min()
                global_max = ref_video.max()
            else:
                global_min = video.min()
                global_max = video.max()

            if vmin is None: vmin = global_min
            if vmax is None: vmax = global_max

        imshow_kwargs = {'cmap': cmap, 'vmin': vmin, 'vmax': vmax}

        frame_idx = 0
        for c in range(ncols):
            if c == spacer_col:
                for r in range(nrows):
                    axes[r, c].axis('off')
                continue
                
            pred_frame = video[frame_idx]
            if rescale: pred_frame = (pred_frame + 1.0) / 2.0
            
            # Fix: Only clip if RGB. Let imshow handle scalar arrays natively.
            if pred_frame.shape[-1] in [3, 4]:
                pred_frame = np.clip(pred_frame, 0.0, 1.0)
            elif pred_frame.shape[-1] == 1:
                pred_frame = pred_frame[..., 0]

            if plot_ref:
                ref_idx = min(frame_idx, ref_video.shape[0] - 1)
                ref_frame = ref_video[ref_idx]
                if rescale: ref_frame = (ref_frame + 1.0) / 2.0
                
                if ref_frame.shape[-1] in [3, 4]:
                    ref_frame = np.clip(ref_frame, 0.0, 1.0)
                elif ref_frame.shape[-1] == 1:
                    ref_frame = ref_frame[..., 0]

            if plot_ref:
                im_ref  = axes[0, c].imshow(ref_frame,  **imshow_kwargs)
                im_pred = axes[1, c].imshow(pred_frame, **imshow_kwargs)
                target_axes = [(axes[0, c], im_ref, ref_frame), (axes[1, c], im_pred, pred_frame)]
            else:
                im_pred = axes[0, c].imshow(pred_frame, **imshow_kwargs)
                target_axes = [(axes[0, c], im_pred, pred_frame)]

            for ax, im_obj, frame_data in target_axes:
                ax.set_xticks([])
                ax.set_yticks([])
                
                h, w = frame_data.shape[:2]
                for spine in ax.spines.values():
                    spine.set_visible(False)
                    
                if show_borders:
                    rect = patches.FancyBboxPatch(
                        (-0.5, -0.5), w, h,
                        boxstyle=f"round,pad=0,rounding_size={corner_radius}", 
                        linewidth=1.2, edgecolor='black', facecolor='none',
                        transform=ax.transData
                    )
                    ax.add_patch(rect)
                    im_obj.set_clip_path(rect)

            if show_titles:
                top_ax = axes[0, c]
                title_str = f"$t={frame_idx + 1}$" if (frame_idx == 0 or (frame_idx + 1 == forecast_start)) else str(frame_idx + 1)
                font_weight = 'bold' if (has_gap and frame_idx + 1 == forecast_start) else 'normal'
                top_ax.set_title(title_str, pad=8, fontsize=18, fontweight=font_weight)

            frame_idx += 1

        if show_labels:
            if plot_ref:
                axes[0, 0].set_ylabel("GT",   rotation=0, labelpad=25, ha='right', va='center', fontsize=28, fontweight='bold')
            axes[-1, 0].set_ylabel("Pred", rotation=0, labelpad=25, ha='right', va='center', fontsize=28, fontweight='bold')

        if save_name:
            plt.savefig(save_name, dpi=100, bbox_inches='tight', facecolor='white', transparent=False)
        else:
            plt.draw()

        try:
            from IPython.display import display
            display(fig)
        except ImportError:
            plt.show()
            
        plt.close(fig)

        # ---------------------------------------------------------
        # GIF Generation
        # ---------------------------------------------------------
        if save_video and save_name is not None:
            
            # Scale fonts up based on gif_scale
            try:
                font = ImageFont.truetype("arial.ttf", 14 * gif_scale)
            except IOError:
                try:
                    font = ImageFont.truetype("DejaVuSans-Bold.ttf", 14 * gif_scale)
                except IOError:
                    font = ImageFont.load_default()

            def process_pil_image(img_array, radius=corner_radius, apply_frame=show_borders):
                h, w = img_array.shape[:2]
                img = Image.fromarray((img_array * 255).astype(np.uint8))
                
                # Resize the underlying image using NEAREST to maintain sharp grid pixels
                new_w, new_h = w * gif_scale, h * gif_scale
                img = img.resize((new_w, new_h), Image.NEAREST)
                
                if not apply_frame: return img
                
                scaled_radius = radius * gif_scale
                mask = Image.new("L", (new_w, new_h), 0)
                draw = ImageDraw.Draw(mask)
                draw.rounded_rectangle((0, 0, new_w, new_h), radius=scaled_radius, fill=255)
                
                rounded_img = Image.new("RGB", (new_w, new_h), "white")
                rounded_img.paste(img, (0, 0), mask=mask)
                
                draw_border = ImageDraw.Draw(rounded_img)
                draw_border.rounded_rectangle((0, 0, new_w-1, new_h-1), radius=scaled_radius, outline="black", width=max(1, gif_scale//2))
                return rounded_img

            def apply_cmap_to_frame(frame, v_min, v_max):
                # Fix: If the frame is already RGB/RGBA, do NOT apply a colormap
                if frame.ndim == 3 and frame.shape[-1] in [3, 4]:
                    # Clip to [0, 1] to be safe, then return the RGB channels
                    return np.clip(frame[..., :3], 0.0, 1.0)
                    
                # If it's a single channel, squeeze it for the colormap
                if frame.ndim == 3 and frame.shape[-1] == 1: 
                    frame = frame[..., 0]
                    
                norm = plt.Normalize(vmin=v_min, vmax=v_max)
                colormap = plt.get_cmap(cmap)
                return colormap(norm(frame))[..., :3]

            gif_frames = []
            scaled_gap = video_gap * gif_scale
            header_height = 20 * gif_scale

            for t in range(nb_frames):
                p_f = video[t]
                if rescale: p_f = (p_f + 1.0) / 2.0
                p_f = apply_cmap_to_frame(p_f, vmin, vmax)

                if plot_ref:
                    r_idx = min(t, ref_video.shape[0] - 1)
                    r_f = ref_video[r_idx]
                    if rescale: r_f = (r_f + 1.0) / 2.0
                    r_f = apply_cmap_to_frame(r_f, vmin, vmax)
                    
                    img_ref  = process_pil_image(r_f)
                    img_pred = process_pil_image(p_f)
                    
                    combined_w = img_ref.width + scaled_gap + img_pred.width
                    combined_h = max(img_ref.height, img_pred.height)
                    combined_frame = Image.new('RGB', (combined_w, combined_h), 'white')
                    combined_frame.paste(img_ref,  (0, 0))
                    combined_frame.paste(img_pred, (img_ref.width + scaled_gap, 0))
                else:
                    combined_frame = process_pil_image(p_f)

                final_img = Image.new('RGB', (combined_frame.width, combined_frame.height + header_height), color='white')
                final_img.paste(combined_frame, (0, header_height))
                
                draw = ImageDraw.Draw(final_img)
                
                if plot_ref:
                    gt_w = draw.textlength("GT", font=font) if hasattr(draw, 'textlength') else 20 * gif_scale
                    pred_w = draw.textlength("Pred", font=font) if hasattr(draw, 'textlength') else 30 * gif_scale
                    
                    draw.text(((img_ref.width - gt_w) // 2, 2 * gif_scale), "GT", font=font, fill="black")
                    draw.text((img_ref.width + scaled_gap + (img_pred.width - pred_w) // 2, 2 * gif_scale), "Pred", font=font, fill="black")
                else:
                    pred_w = draw.textlength("Pred", font=font) if hasattr(draw, 'textlength') else 30 * gif_scale
                    draw.text(((combined_frame.width - pred_w) // 2, 2 * gif_scale), "Pred", font=font, fill="black")
                
                gif_frames.append(final_img)
            
            gif_path = Path(save_name).with_suffix('.gif')
            gif_frames[0].save(gif_path, save_all=True, append_images=gif_frames[1:], duration=150, loop=0)
            print(f"Saved rollout animation to {gif_path}")

            try:
                from IPython.display import Image as IPyImage, display
                display(IPyImage(filename=str(gif_path)))
            except ImportError:
                pass