tasks/image_classification/plotting.py

import numpy as np
import cv2
import torch
import os
import imageio
import matplotlib.pyplot as plt
import matplotlib as mpl
from matplotlib import patheffects
mpl.use('Agg')
import seaborn as sns
import numpy as np
from tqdm.auto import tqdm
sns.set_style('darkgrid')

from tqdm.auto import tqdm
from scipy import ndimage
import umap
from scipy.special import softmax

import subprocess as sp
import cv2 # Still potentially useful for color conversion checks if needed
import os

def save_frames_to_mp4(frames, output_filename, fps=15.0, gop_size=None, crf=23, preset='medium', pix_fmt='yuv420p'):
    """
    Saves a list of NumPy array frames to an MP4 video file using FFmpeg via subprocess.

    Includes fix for odd frame dimensions by padding to the nearest even number using -vf pad.

    Requires FFmpeg to be installed and available in the system PATH.

    Args:
        frames (list): A list of NumPy arrays representing the video frames.
                       Expected format: uint8, (height, width, 3) for BGR color
                       or (height, width) for grayscale. Should be consistent.
        output_filename (str): The path and name for the output MP4 file.
        fps (float, optional): Frames per second for the output video. Defaults to 15.0.
        gop_size (int, optional): Group of Pictures (GOP) size. This determines the
                                  maximum interval between keyframes. Lower values
                                  mean more frequent keyframes (better seeking, larger file).
                                  Defaults to int(fps) (approx 1 keyframe per second).
        crf (int, optional): Constant Rate Factor for H.264 encoding. Lower values mean
                             better quality and larger files. Typical range: 18-28.
                             Defaults to 23.
        preset (str, optional): FFmpeg encoding speed preset. Affects encoding time
                                and compression efficiency. Options include 'ultrafast',
                                'superfast', 'veryfast', 'faster', 'fast', 'medium',
                                'slow', 'slower', 'veryslow'. Defaults to 'medium'.
    """
    if not frames:
        print("Error: The 'frames' list is empty. No video to save.")
        return

    # --- Determine Parameters from First Frame ---
    try:
        first_frame = frames[0]
        print(first_frame.shape)
        if not isinstance(first_frame, np.ndarray):
             print(f"Error: Frame 0 is not a NumPy array (type: {type(first_frame)}).")
             return

        frame_height, frame_width = first_frame.shape[:2]
        frame_size_str = f"{frame_width}x{frame_height}"

        # Determine input pixel format based on first frame's shape
        if len(first_frame.shape) == 3 and first_frame.shape[2] == 3:
            input_pixel_format = 'bgr24' # Assume OpenCV's default BGR uint8
            expected_dims = 3
            print(f"Info: Detected color frames (shape: {first_frame.shape}). Expecting BGR input.")
        elif len(first_frame.shape) == 2:
            input_pixel_format = 'gray'
            expected_dims = 2
            print(f"Info: Detected grayscale frames (shape: {first_frame.shape}).")
        else:
            print(f"Error: Unsupported frame shape {first_frame.shape}. Must be (h, w) or (h, w, 3).")
            return

        if first_frame.dtype != np.uint8:
             print(f"Warning: First frame dtype is {first_frame.dtype}. Will attempt conversion to uint8.")

    except IndexError:
        print("Error: Could not access the first frame to determine dimensions.")
        return
    except Exception as e:
         print(f"Error processing first frame: {e}")
         return

    # --- Set GOP size default if not provided ---
    if gop_size is None:
        gop_size = int(fps)
        print(f"Info: GOP size not specified, defaulting to {gop_size} (approx 1 keyframe/sec).")

    # --- Construct FFmpeg Command ---
    # ADDED -vf pad filter to ensure even dimensions for libx264/yuv420p
    # It calculates the nearest even dimensions >= original dimensions
    # Example: 1600x1351 -> 1600x1352
    pad_filter = "pad=ceil(iw/2)*2:ceil(ih/2)*2"

    command = [
        'ffmpeg',
        '-y',
        '-f', 'rawvideo',
        '-vcodec', 'rawvideo',
        '-pix_fmt', input_pixel_format,
        '-s', frame_size_str,
        '-r', str(float(fps)),
        '-i', '-',
        '-vf', pad_filter, # <--- ADDED VIDEO FILTER HERE
        '-c:v', 'libx264',
        '-pix_fmt', pix_fmt,
        '-preset', preset,
        '-crf', str(crf),
        '-g', str(gop_size),
        '-movflags', '+faststart',
        output_filename
    ]

    print(f"\n--- Starting FFmpeg ---")
    print(f"Output File: {output_filename}")
    print(f"Parameters: FPS={fps}, Size={frame_size_str}, GOP={gop_size}, CRF={crf}, Preset={preset}")
    print(f"Applying Filter: -vf {pad_filter} (Ensures even dimensions)")
    # print(f"FFmpeg Command: {' '.join(command)}") # Uncomment for debugging

    # --- Execute FFmpeg via Subprocess ---
    try:
        process = sp.Popen(command, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE)

        print(f"\nWriting {len(frames)} frames to FFmpeg...")
        progress_interval = max(1, len(frames) // 10) # Print progress roughly 10 times

        for i, frame in enumerate(frames):
            # Basic validation and conversion for each frame
            if not isinstance(frame, np.ndarray):
                 print(f"Warning: Frame {i} is not a numpy array (type: {type(frame)}). Skipping.")
                 continue
            if frame.shape[0] != frame_height or frame.shape[1] != frame_width:
                print(f"Warning: Frame {i} has different dimensions {frame.shape[:2]}! Expected ({frame_height},{frame_width}). Skipping.")
                continue

            current_dims = len(frame.shape)
            if current_dims != expected_dims:
                 print(f"Warning: Frame {i} has inconsistent dimensions ({current_dims}D vs expected {expected_dims}D). Skipping.")
                 continue
            if expected_dims == 3 and frame.shape[2] != 3:
                 print(f"Warning: Frame {i} is color but doesn't have 3 channels ({frame.shape}). Skipping.")
                 continue

            if frame.dtype != np.uint8:
                try:
                    frame = np.clip(frame, 0, 255).astype(np.uint8)
                except Exception as clip_err:
                     print(f"Error clipping/converting frame {i} dtype: {clip_err}. Skipping.")
                     continue

            # Write frame bytes to FFmpeg's stdin
            try:
                 process.stdin.write(frame.tobytes())
            except (OSError, BrokenPipeError) as pipe_err:
                 print(f"\nError writing frame {i} to FFmpeg stdin: {pipe_err}")
                 print("FFmpeg process likely terminated prematurely. Check FFmpeg errors below.")
                 try:
                     # Immediately try to read stderr if pipe breaks
                     stderr_output_on_error = process.stderr.read()
                     if stderr_output_on_error:
                          print("\n--- FFmpeg stderr output on error ---")
                          print(stderr_output_on_error.decode(errors='ignore'))
                          print("--- End FFmpeg stderr ---")
                 except Exception as read_err:
                     print(f"(Could not read stderr after pipe error: {read_err})")
                 return
            except Exception as write_err:
                 print(f"Unexpected error writing frame {i}: {write_err}. Skipping.")
                 continue

            if (i + 1) % progress_interval == 0 or (i + 1) == len(frames):
                 print(f"  Processed frame {i + 1}/{len(frames)}")

        print("\nFinished writing frames. Closing FFmpeg stdin and waiting for completion...")
        process.stdin.close()
        stdout, stderr = process.communicate()
        return_code = process.wait()

        print("\n--- FFmpeg Final Status ---")
        if return_code == 0:
            print(f"FFmpeg process completed successfully.")
            print(f"Video saved as: {output_filename}")
        else:
            print(f"FFmpeg process failed with return code {return_code}.")
            print("--- FFmpeg Standard Error Output: ---")
            print(stderr.decode(errors='replace')) # Print stderr captured by communicate()
            print("--- End FFmpeg Output ---")
            print("Review the FFmpeg error message above for details (e.g., dimension errors, parameter issues).")

    except FileNotFoundError:
        print("\n--- FATAL ERROR ---")
        print("Error: 'ffmpeg' command not found.")
        print("Please ensure FFmpeg is installed and its directory is included in your system's PATH environment variable.")
        print("Download from: https://ffmpeg.org/")
        print("-------------------")
    except Exception as e:
        print(f"\nAn unexpected error occurred during FFmpeg execution: {e}")

def find_island_centers(array_2d, threshold):
    """
    Finds the center of mass of each island (connected component) in a 2D array.

    Args:
        array_2d: A 2D numpy array of values.
        threshold: The threshold to binarize the array.

    Returns:
        A list of tuples (y, x) representing the center of mass of each island.
    """
    binary_image = array_2d > threshold
    labeled_image, num_labels = ndimage.label(binary_image)
    centers = []
    areas = []  # Store the area of each island
    for i in range(1, num_labels + 1):
        island = (labeled_image == i)
        total_mass = np.sum(array_2d[island])
        if total_mass > 0:
            y_coords, x_coords = np.mgrid[:array_2d.shape[0], :array_2d.shape[1]]
            x_center = np.average(x_coords[island], weights=array_2d[island])
            y_center = np.average(y_coords[island], weights=array_2d[island])
            centers.append((round(y_center, 4), round(x_center, 4)))
            areas.append(np.sum(island))  # Calculate area of the island
    return centers, areas

def plot_neural_dynamics(post_activations_history, N_to_plot, save_location, axis_snap=False, N_per_row=5, which_neurons_mid=None, mid_colours=None, use_most_active_neurons=False):
    assert N_to_plot%N_per_row==0, f'For nice visualisation, N_to_plot={N_to_plot} must be a multiple of N_per_row={N_per_row}'
    assert post_activations_history.shape[-1] >= N_to_plot
    figscale = 2
    aspect_ratio = 3
    mosaic = np.array([[f'{i}'] for i in range(N_to_plot)]).flatten().reshape(-1, N_per_row)
    fig_synch, axes_synch = plt.subplot_mosaic(mosaic=mosaic, figsize=(figscale*mosaic.shape[1]*aspect_ratio*0.2, figscale*mosaic.shape[0]*0.2))
    fig_mid, axes_mid = plt.subplot_mosaic(mosaic=mosaic, figsize=(figscale*mosaic.shape[1]*aspect_ratio*0.2, figscale*mosaic.shape[0]*0.2), dpi=200)

    palette = sns.color_palette("husl", 8)
    
    which_neurons_synch = np.arange(N_to_plot)
    # which_neurons_mid = np.arange(N_to_plot, N_to_plot*2) if post_activations_history.shape[-1] >= 2*N_to_plot else np.random.choice(np.arange(post_activations_history.shape[-1]), size=N_to_plot, replace=True)
    random_indices = np.random.choice(np.arange(post_activations_history.shape[-1]), size=N_to_plot, replace=post_activations_history.shape[-1] < N_to_plot)
    if use_most_active_neurons:
        metric = np.abs(np.fft.rfft(post_activations_history, axis=0))[3:].mean(0).std(0)
        random_indices = np.argsort(metric)[-N_to_plot:]
        np.random.shuffle(random_indices)
    which_neurons_mid = which_neurons_mid if which_neurons_mid is not None else random_indices

    if mid_colours is None:
        mid_colours = [palette[np.random.randint(0, 8)] for ndx in range(N_to_plot)]
    with tqdm(total=N_to_plot, initial=0, leave=False, position=1, dynamic_ncols=True) as pbar_inner:
        pbar_inner.set_description('Plotting neural dynamics')
        for ndx in range(N_to_plot):
            
            ax_s = axes_synch[f'{ndx}']
            ax_m = axes_mid[f'{ndx}']

            traces_s = post_activations_history[:,:,which_neurons_synch[ndx]].T
            traces_m = post_activations_history[:,:,which_neurons_mid[ndx]].T
            c_s = palette[np.random.randint(0, 8)]
            c_m = mid_colours[ndx]

            for traces_s_here, traces_m_here in zip(traces_s, traces_m):
                ax_s.plot(np.arange(len(traces_s_here)), traces_s_here, linestyle='-', color=c_s, alpha=0.05, linewidth=0.6)
                ax_m.plot(np.arange(len(traces_m_here)), traces_m_here, linestyle='-', color=c_m, alpha=0.05, linewidth=0.6)

            
            ax_s.plot(np.arange(len(traces_s[0])), traces_s[0], linestyle='-', color='white', alpha=1, linewidth=2.5)
            ax_s.plot(np.arange(len(traces_s[0])), traces_s[0], linestyle='-', color=c_s, alpha=1, linewidth=1.3)
            ax_s.plot(np.arange(len(traces_s[0])), traces_s[0], linestyle='-', color='black', alpha=1, linewidth=0.3)
            ax_m.plot(np.arange(len(traces_m[0])), traces_m[0], linestyle='-', color='white', alpha=1, linewidth=2.5)
            ax_m.plot(np.arange(len(traces_m[0])), traces_m[0], linestyle='-', color=c_m, alpha=1, linewidth=1.3)
            ax_m.plot(np.arange(len(traces_m[0])), traces_m[0], linestyle='-', color='black', alpha=1, linewidth=0.3)
            if axis_snap and np.all(np.isfinite(traces_s[0])):
                ax_s.set_ylim([np.min(traces_s[0])-np.ptp(traces_s[0])*0.05, np.max(traces_s[0])+np.ptp(traces_s[0])*0.05])
                ax_m.set_ylim([np.min(traces_m[0])-np.ptp(traces_m[0])*0.05, np.max(traces_m[0])+np.ptp(traces_m[0])*0.05])
            

            ax_s.grid(False)
            ax_m.grid(False)
            ax_s.set_xlim([0, len(traces_s[0])-1])
            ax_m.set_xlim([0, len(traces_m[0])-1])

            ax_s.set_xticklabels([])
            ax_s.set_yticklabels([])

            ax_m.set_xticklabels([])
            ax_m.set_yticklabels([])
            pbar_inner.update(1)
    fig_synch.tight_layout(pad=0.05)
    fig_mid.tight_layout(pad=0.05)
    if save_location is not None:
        fig_synch.savefig(f'{save_location}/neural_dynamics_synch.pdf', dpi=200)
        fig_synch.savefig(f'{save_location}/neural_dynamics_synch.png', dpi=200)
        fig_mid.savefig(f'{save_location}/neural_dynamics_other.pdf', dpi=200)
        fig_mid.savefig(f'{save_location}/neural_dynamics_other.png', dpi=200)
        plt.close(fig_synch)
        plt.close(fig_mid)
    return fig_synch, fig_mid, which_neurons_mid, mid_colours



def make_classification_gif(image, target, predictions, certainties, post_activations, attention_tracking, class_labels, save_location):
    cmap_viridis = sns.color_palette('viridis', as_cmap=True)
    cmap_spectral = sns.color_palette("Spectral", as_cmap=True)
    figscale = 2
    with tqdm(total=post_activations.shape[0]+1, initial=0, leave=False, position=1, dynamic_ncols=True) as pbar_inner:
        pbar_inner.set_description('Computing UMAP')
    

        low = np.percentile(post_activations, 1, axis=0, keepdims=True)
        high = np.percentile(post_activations, 99, axis=0, keepdims=True)
        post_activations_normed = np.clip((post_activations - low)/(high - low), 0, 1)
        metric = 'cosine'
        reducer = umap.UMAP(n_components=2,
                            n_neighbors=100,
                            min_dist=3,
                            spread=3.0,
                            metric=metric,
                            random_state=None,
                            # low_memory=True,
                            ) if post_activations.shape[-1] > 2048 else umap.UMAP(n_components=2,
                            n_neighbors=20,
                            min_dist=1,
                            spread=1.0,
                            metric=metric,
                            random_state=None,
                            # low_memory=True,
                            )
        positions = reducer.fit_transform(post_activations_normed.T)

        x_umap = positions[:, 0]
        y_umap = positions[:, 1]

        pbar_inner.update(1)
        pbar_inner.set_description('Iterating through to build frames')


        
        frames = []
        route_steps = {}
        route_colours = []

        n_steps = len(post_activations)
        n_heads = attention_tracking.shape[1]
        step_linspace = np.linspace(0, 1, n_steps)
        
        for stepi in np.arange(0, n_steps, 1):
            pbar_inner.set_description('Making frames for gif')
            

            attention_now = attention_tracking[max(0, stepi-5):stepi+1].mean(0)  # Make it smooth for pretty
            # attention_now[:,0,0] = 0  # Corners can be weird looking
            # attention_now[:,0,-1] = 0
            # attention_now[:,-1,0] = 0
            # attention_now[:,-1,-1] = 0
            # attention_now = (attention_tracking[:stepi+1, 0] * decay).sum(0)/(decay.sum(0))
            certainties_now = certainties[1, :stepi+1]
            attention_interp = torch.nn.functional.interpolate(torch.from_numpy(attention_now).unsqueeze(0), image.shape[:2], mode='bilinear')[0]
            attention_interp = (attention_interp.flatten(1) - attention_interp.flatten(1).min(-1, keepdim=True)[0])/(attention_interp.flatten(1).max(-1, keepdim=True)[0] - attention_interp.flatten(1).min(-1, keepdim=True)[0])
            attention_interp = attention_interp.reshape(n_heads, image.shape[0], image.shape[1])
            

            colour = list(cmap_spectral(step_linspace[stepi]))
            route_colours.append(colour)
            for headi in range(min(8, n_heads)):
                com_attn = np.copy(attention_interp[headi])
                com_attn[com_attn < np.percentile(com_attn, 97)] = 0.0
                if headi not in route_steps:
                    A = attention_interp[headi].detach().cpu().numpy()
                    centres, areas = find_island_centers(A, threshold=0.7)
                    route_steps[headi] = [centres[np.argmax(areas)]]
                else:
                    A = attention_interp[headi].detach().cpu().numpy()
                    centres, areas = find_island_centers(A, threshold=0.7)
                    route_steps[headi] = route_steps[headi] + [centres[np.argmax(areas)]]

            mosaic = [['head_0', 'head_0_overlay', 'head_1', 'head_1_overlay'],
                      ['head_2', 'head_2_overlay', 'head_3', 'head_3_overlay'],
                      ['head_4', 'head_4_overlay', 'head_5', 'head_5_overlay'],
                      ['head_6', 'head_6_overlay', 'head_7', 'head_7_overlay'],
                      ['probabilities', 'probabilities','certainty', 'certainty'],
                      ['umap', 'umap', 'umap', 'umap'],
                      ['umap', 'umap', 'umap', 'umap'],
                      ['umap', 'umap', 'umap', 'umap'],
                     
                      ]
            
            
            img_aspect = image.shape[0]/image.shape[1]
            # print(img_aspect)
            aspect_ratio = (4*figscale, 8*figscale*img_aspect)
            fig, axes = plt.subplot_mosaic(mosaic, figsize=aspect_ratio)
            for ax in axes.values():
                ax.axis('off')


            axes['certainty'].plot(np.arange(len(certainties_now)), certainties_now, 'k-', linewidth=figscale*1, label='1-(normalised entropy)')
            for ii, (x, y) in enumerate(zip(np.arange(len(certainties_now)), certainties_now)):
                is_correct = predictions[:, ii].argmax(-1)==target
                if is_correct: axes['certainty'].axvspan(ii, ii + 1, facecolor='limegreen', edgecolor=None, lw=0, alpha=0.3)
                else:
                    axes['certainty'].axvspan(ii, ii + 1, facecolor='orchid', edgecolor=None, lw=0, alpha=0.3)
            axes['certainty'].plot(len(certainties_now)-1, certainties_now[-1], 'k.', markersize=figscale*4)
            axes['certainty'].axis('off')
            axes['certainty'].set_ylim([-0.05, 1.05])
            axes['certainty'].set_xlim([0, certainties.shape[-1]+1])

            ps = torch.softmax(torch.from_numpy(predictions[:, stepi]), -1)
            k = 15 if len(class_labels) > 15 else len(class_labels)
            topk = torch.topk (ps, k, dim = 0, largest=True).indices.detach().cpu().numpy()
            top_classes = np.array(class_labels)[topk]
            true_class = target
            colours = [('b' if ci != true_class else 'g') for ci in topk]
            bar_heights = ps[topk].detach().cpu().numpy()


            axes['probabilities'].bar(np.arange(len(bar_heights))[::-1], bar_heights, color=np.array(colours), alpha=1)
            axes['probabilities'].set_ylim([0, 1])
            axes['probabilities'].axis('off')


            for i, (name) in enumerate(top_classes):
                prob = ps[i]
                is_correct = name==class_labels[true_class]
                fg_color = 'darkgreen' if is_correct else 'crimson'
                text_str = f'{name[:40]}'
                axes['probabilities'].text(
                    0.05,
                    0.95 - i * 0.055,  # Adjust vertical position for each line
                    text_str,
                    transform=axes['probabilities'].transAxes,
                    verticalalignment='top',
                    fontsize=8,  # Increased font size
                    color=fg_color,
                    alpha=0.5,
                    path_effects=[
                        patheffects.Stroke(linewidth=3, foreground='aliceblue'),
                        patheffects.Normal()
                    ])
            


            attention_now = attention_tracking[max(0, stepi-5):stepi+1].mean(0)  # Make it smooth for pretty
            # attention_now = (attention_tracking[:stepi+1, 0] * decay).sum(0)/(decay.sum(0))
            certainties_now = certainties[1, :stepi+1]
            attention_interp = torch.nn.functional.interpolate(torch.from_numpy(attention_now).unsqueeze(0), image.shape[:2], mode='nearest')[0]
            attention_interp = (attention_interp.flatten(1) - attention_interp.flatten(1).min(-1, keepdim=True)[0])/(attention_interp.flatten(1).max(-1, keepdim=True)[0] - attention_interp.flatten(1).min(-1, keepdim=True)[0])
            attention_interp = attention_interp.reshape(n_heads, image.shape[0], image.shape[1])

            for hi in range(min(8, n_heads)):
                ax = axes[f'head_{hi}']
                img_to_plot = cmap_viridis(attention_interp[hi].detach().cpu().numpy())
                ax.imshow(img_to_plot)

                ax_overlay = axes[f'head_{hi}_overlay']

                these_route_steps = route_steps[hi]
                y_coords, x_coords = zip(*these_route_steps)
                y_coords = image.shape[-2] - np.array(list(y_coords))-1
                
                ax_overlay.imshow(np.flip(image, axis=0), origin='lower')
                # ax.imshow(np.flip(solution_maze, axis=0), origin='lower')
                arrow_scale = 1.5 if image.shape[0] > 32 else 0.8
                for i in range(len(these_route_steps)-1):
                    dx = x_coords[i+1] - x_coords[i]
                    dy = y_coords[i+1] - y_coords[i]

                    ax_overlay.arrow(x_coords[i], y_coords[i], dx, dy, linewidth=1.6*arrow_scale*1.3, head_width=1.9*arrow_scale*1.3, head_length=1.4*arrow_scale*1.45, fc='white', ec='white', length_includes_head = True, alpha=1)
                    ax_overlay.arrow(x_coords[i], y_coords[i], dx, dy, linewidth=1.6*arrow_scale, head_width=1.9*arrow_scale, head_length=1.4*arrow_scale, fc=route_colours[i], ec=route_colours[i], length_includes_head = True)
                    
                ax_overlay.set_xlim([0,image.shape[1]-1])
                ax_overlay.set_ylim([0,image.shape[0]-1])
                ax_overlay.axis('off')


            z = post_activations_normed[stepi]

            axes['umap'].scatter(x_umap, y_umap, s=30, c=cmap_spectral(z))
        
            fig.tight_layout(pad=0.1)
            


            canvas = fig.canvas
            canvas.draw()
            image_numpy = np.frombuffer(canvas.buffer_rgba(), dtype='uint8')
            image_numpy = (image_numpy.reshape(*reversed(canvas.get_width_height()), 4)[:,:,:3])
            frames.append(image_numpy)
            plt.close(fig)
            pbar_inner.update(1)
        pbar_inner.set_description('Saving gif')
        imageio.mimsave(save_location, frames, fps=15, loop=100)