import numpy as np
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes import Axes

def visualize_vector(ax: Axes, vector: np.ndarray) -> None:
    """
    Visualizes a 1D vector on the given Matplotlib axes.

    Args:
        ax: The Matplotlib axes object to draw on.
        vector: The 1D NumPy array to plot.
    """
    ax.clear()
    ax.plot(vector)
    ax.set_title("Vector Data")
    ax.set_xlabel("Index")
    ax.set_ylabel("Value")
    ax.grid(True)
    
    # Set Y-axis limits to fit the data, with a small padding
    data_min = np.min(vector)
    data_max = np.max(vector)
    padding = (data_max - data_min) * 0.05  # 5% padding
    
    # Handle the case of a flat line where min == max
    if padding == 0:
        padding = 1.0

    ax.set_ylim(data_min - padding, data_max + padding)


def visualize_matrix(
    fig: Figure, 
    ax: Axes, 
    matrix: np.ndarray, 
    colormap: str = 'viridis'
) -> None:
    """
    Visualizes a 2D matrix as an image on the given Matplotlib axes.

    Args:
        fig: The Matplotlib figure, needed for the colorbar.
        ax: The Matplotlib axes object to draw on.
        matrix: The 2D NumPy array to visualize.
        colormap: The name of the colormap to use.
    """
    ax.clear()
    im = ax.imshow(matrix, cmap=colormap, interpolation='nearest', origin='lower')
    
    ax.set_title("Matrix Data")
    ax.set_xlabel("Column")
    ax.set_ylabel("Row")
    
    # Add or update the colorbar
    # We remove any existing colorbar first to prevent duplicates
    if fig.axes[-1] is not ax: # A simple check to see if a colorbar might already exist
         if len(fig.axes) > 1 and fig.axes[-1] is not ax:
            fig.delaxes(fig.axes[-1])

    fig.colorbar(im, ax=ax, label='Value')