from ._sharedutils import generate_abc, load_devanagari_font

import matplotlib.pyplot as plt
import numpy as np

def tracks(
    config,
    ylims=None,
    xlim=None,
    ylabels=None,
    xlabels=None,
    titles=None,
    grid=True,
    fig_width=18,
    hspace_inches=0.4,
    abc=True,
    fig_num="",
    ):
    """
    Create a Praat-style multi-track figure layout where all tracks share the same x-axis 
    (typically representing time), with consistent physical heights and uniform spacing 
    between tracks.

    This method provides a flexible and publication-ready layout generator for 
    time-aligned visualizations such as signals, annotations, and spectrograms.

    Parameters
    ----------
    config : str
        A string specifying the types and order of tracks to display.
        Each character represents one track:
        - "a": auxiliary track (e.g., annotation, pitch, energy, etc.)
        - "s": signal track (e.g., waveform)
        - "m": matrix track (e.g., spectrogram, feature matrix)
        Example:
            "asm" → one auxiliary track, one signal track, one matrix track.
    ylims : list[tuple[float, float]] or None, optional
        A list of (min, max) pairs specifying y-axis limits for each track.
        Length must match the number of tracks if provided.
    xlim : tuple[float, float] or None, optional
        The x-axis limit shared across all tracks.
    ylabels : list[str] or None, optional
        A list of y-axis labels for each track.
    xlabels : list[str] or None, optional
        A list of x-axis labels for each track.
    titles : list[str] or None, optional
        A list of titles for each track.
    grid: bool, default=True
        Adds gridlines to the tracks.
    fig_width : float, default=18
        Total figure width (in inches).
    hspace_inches : float, default=0.4
        Vertical gap (in inches) between consecutive tracks.
        This is converted internally to a relative spacing value to preserve
        exact physical spacing across configurations.
    abc : bool, default=True
        Whether to label each subplot with a small reference tag (e.g., "(a)", "(b)", "(c)")
        for research paper style referencing.
    fig_num : str or int or float, default=""
        Optional prefix to the alphabetical subplot labels, e.g. "1a", "2b", etc.
    
    Returns
    -------
    fig : matplotlib.figure.Figure
        The created Matplotlib Figure object.
    axs : numpy.ndarray of matplotlib.axes.Axes
        A 2D array of Axes objects of shape (n_tracks, 2), where column 0 contains
        the actual tracks and column 1 is reserved for colorbars (for matrix plots).
    """

    load_devanagari_font()
    
    # Setup the variables
    tracks = config
    
    # Parse the config
    n_aux_tracks = config.count("a")
    n_signal_tracks = config.count("s")
    n_matrix_tracks = config.count("m")
    n_tracks = n_aux_tracks + n_signal_tracks + n_matrix_tracks  # Number of tracks
    
    # Decide heights of different tracks type
    height = {}
    height["a"] = 0.4  # Aux height
    height["s"] = 2.0  # Signal height
    height["m"] = 4.0  # Matrix height
    cbar_width = 0.01  # For second column (for matrix track)
    
    # Calculate height ratios for each track
    height_ratios = [height[track] for track in tracks]
        
    # Calculate total fig height
    fig_height = np.sum(height_ratios) + (len(tracks) - 1) * hspace_inches
    
    # Compute hspace wrt the average subplot heights (matplot uses this)
    total_subplot_height = sum(height_ratios)
    avg_subplot_height = total_subplot_height / len(tracks)
    hspace = hspace_inches / avg_subplot_height
    
    # Generate the layout
    fig, axs = plt.subplots(n_tracks, 2, figsize=(fig_width, fig_height), gridspec_kw={"height_ratios": height_ratios, "width_ratios": [1, cbar_width], "hspace": hspace})
    
    axs = np.atleast_2d(axs)  # This is done for consistency otherwise axs[i, 0] does not work

    for i, track in enumerate(tracks):  # Loop through each tier and adjust the layout
        if track == "a":  # Remove ticks and labels from all the aux subplots
            axs[i, 0].tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
        elif track == "s":
            axs[i, 0].tick_params(bottom=False, labelbottom=False)
        elif track == "m":
            axs[i, 0].tick_params(bottom=False, labelbottom=False)
        
        axs[i, 1].axis("off")  # Turn off the column 2, only turn it on when matrix is plotted with colorbar
        
        axs[i, 0].sharex(axs[0, 0])  # Share the x-axis to make all the tiers aligned
    
    # Add tags (1a, 4.1c, ...) to each tier for better referencing in research papers.
    if abc is True:
        abc_labels = generate_abc(n_tracks)
        for i in range(n_tracks):
            axs[i, 0].set_title(f"({fig_num}{abc_labels[i]})", fontsize=10, loc="left")
            
    # Turn on the x-label for the last tier
    axs[-1, 0].tick_params(bottom=True, labelbottom=True)
    
    # Configure the styling
    if ylims is not None:
        for (ylim, ax) in zip(ylims, axs[:,0]):
            set.limit(ax, ylim=ylim)
    if xlim is not None:
        set.limit(axs[0,0], xlim=xlim)
    if ylabels is not None:
        for (ylabel, ax) in zip(ylabels, axs[:,0]):
            set.label(ax, ylabel=ylabel)
    if xlabels is not None:
        for (xlabel, ax) in zip(xlabels, axs[:,0]):
            set.label(ax, xlabel=xlabel)
    if titles is not None:
        for (title, ax) in zip(titles, axs[:,0]):
            set.title(ax, title=title)
    if grid is True:
        for ax in axs[:,0]:
            set.gridlines(ax, x=True, y=True)
        
    return fig, axs
