"""
europa.display
~~~~~~~~~~~~~~
Visualization utilities for audio data.

Functions:
    waveform        - Plot audio waveform
    spectrogram     - Plot spectrogram
    mel_spectrogram - Plot mel spectrogram
    mfcc            - Plot MFCCs
    chroma          - Plot chromagram
    beats           - Plot waveform with beat markers
    all_features    - Multi-panel feature overview
"""

import numpy as np
from ..core.audio import Audio


def waveform(
    audio: Audio,
    title: str = "Waveform",
    figsize: tuple = (12, 3),
    color: str = "#4A90D9",
    show: bool = True,
) -> None:
    """
    Plot the audio waveform.

    Args:
        audio   : Input Audio.
        title   : Plot title.
        figsize : Figure size (width, height).
        color   : Waveform color.
        show    : If True, call plt.show().
    """
    import matplotlib.pyplot as plt
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    t = np.linspace(0, audio.duration, len(data))

    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(t, data, color=color, linewidth=0.5)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Amplitude")
    ax.set_title(title)
    ax.set_xlim(0, audio.duration)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def spectrogram(
    audio: Audio,
    n_fft: int = 2048,
    hop_length: int = 512,
    title: str = "Spectrogram",
    figsize: tuple = (12, 4),
    cmap: str = "magma",
    show: bool = True,
) -> None:
    """
    Plot a log-power spectrogram.

    Args:
        audio      : Input Audio.
        n_fft      : FFT size.
        hop_length : Hop size.
        title      : Plot title.
        figsize    : Figure size.
        cmap       : Matplotlib colormap.
        show       : If True, call plt.show().
    """
    import matplotlib.pyplot as plt
    from ..feature import spectrogram as spec_feat

    S = spec_feat(audio, n_fft=n_fft, hop_length=hop_length, power=2.0)
    S_db = 10 * np.log10(S + 1e-9)

    fig, ax = plt.subplots(figsize=figsize)
    img = ax.imshow(
        S_db, aspect="auto", origin="lower", cmap=cmap,
        extent=[0, audio.duration, 0, audio.sr / 2 / 1000],
    )
    plt.colorbar(img, ax=ax, label="Power (dB)")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Frequency (kHz)")
    ax.set_title(title)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def mel_spectrogram(
    audio: Audio,
    n_mels: int = 128,
    hop_length: int = 512,
    title: str = "Mel Spectrogram",
    figsize: tuple = (12, 4),
    cmap: str = "magma",
    show: bool = True,
) -> None:
    """
    Plot a mel-scaled spectrogram.
    """
    import matplotlib.pyplot as plt
    from ..feature import mel_spectrogram as mel_feat

    M = mel_feat(audio, n_mels=n_mels, hop_length=hop_length)
    M_db = 10 * np.log10(M + 1e-9)

    fig, ax = plt.subplots(figsize=figsize)
    img = ax.imshow(
        M_db, aspect="auto", origin="lower", cmap=cmap,
        extent=[0, audio.duration, 0, n_mels],
    )
    plt.colorbar(img, ax=ax, label="Power (dB)")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Mel band")
    ax.set_title(title)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def mfcc(
    audio: Audio,
    n_mfcc: int = 13,
    hop_length: int = 512,
    title: str = "MFCCs",
    figsize: tuple = (12, 4),
    cmap: str = "coolwarm",
    show: bool = True,
) -> None:
    """
    Plot Mel-Frequency Cepstral Coefficients.
    """
    import matplotlib.pyplot as plt
    from ..feature import mfcc as mfcc_feat

    M = mfcc_feat(audio, n_mfcc=n_mfcc, hop_length=hop_length)

    fig, ax = plt.subplots(figsize=figsize)
    img = ax.imshow(
        M, aspect="auto", origin="lower", cmap=cmap,
        extent=[0, audio.duration, 0, n_mfcc],
    )
    plt.colorbar(img, ax=ax, label="Coefficient")
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("MFCC coefficient")
    ax.set_title(title)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def chroma(
    audio: Audio,
    hop_length: int = 512,
    title: str = "Chromagram",
    figsize: tuple = (12, 4),
    cmap: str = "viridis",
    show: bool = True,
) -> None:
    """
    Plot chromagram.
    """
    import matplotlib.pyplot as plt
    from ..feature import chroma as chroma_feat

    C = chroma_feat(audio, hop_length=hop_length)
    note_names = ['C','C#','D','D#','E','F','F#','G','G#','A','A#','B']

    fig, ax = plt.subplots(figsize=figsize)
    img = ax.imshow(C, aspect="auto", origin="lower", cmap=cmap,
                    extent=[0, audio.duration, 0, 12])
    plt.colorbar(img, ax=ax, label="Energy")
    ax.set_yticks(np.arange(12) + 0.5)
    ax.set_yticklabels(note_names)
    ax.set_xlabel("Time (s)")
    ax.set_title(title)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def beats(
    audio: Audio,
    beat_times: np.ndarray = None,
    title: str = "Waveform + Beats",
    figsize: tuple = (12, 3),
    show: bool = True,
) -> None:
    """
    Plot waveform with beat locations overlaid.

    Args:
        audio      : Input Audio.
        beat_times : Beat positions in seconds. If None, runs beat_track.
        title      : Plot title.
        figsize    : Figure size.
        show       : If True, call plt.show().
    """
    import matplotlib.pyplot as plt
    if beat_times is None:
        from ..analysis import beat_track
        _, beat_times = beat_track(audio)

    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    t = np.linspace(0, audio.duration, len(data))

    fig, ax = plt.subplots(figsize=figsize)
    ax.plot(t, data, color="#888", linewidth=0.5, alpha=0.7)
    for bt in beat_times:
        ax.axvline(bt, color="#FF5733", linewidth=0.8, alpha=0.8)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Amplitude")
    ax.set_title(title)
    ax.set_xlim(0, audio.duration)
    plt.tight_layout()
    if show:
        plt.show()
    return fig


def all_features(
    audio: Audio,
    figsize: tuple = (14, 10),
    show: bool = True,
) -> None:
    """
    Multi-panel feature overview: waveform, spectrogram, mel, MFCCs, chroma.
    """
    import matplotlib.pyplot as plt
    from ..feature import spectrogram as spec_feat, mel_spectrogram as mel_feat
    from ..feature import mfcc as mfcc_feat, chroma as chroma_feat

    fig, axes = plt.subplots(5, 1, figsize=figsize)
    data = audio.data if audio.is_mono else np.mean(audio.data, axis=0)
    t = np.linspace(0, audio.duration, len(data))

    # Waveform
    axes[0].plot(t, data, linewidth=0.5, color="#4A90D9")
    axes[0].set_title("Waveform")
    axes[0].set_xlim(0, audio.duration)

    def _show(ax, matrix, label, cmap="magma"):
        ax.imshow(matrix, aspect="auto", origin="lower", cmap=cmap,
                  extent=[0, audio.duration, 0, matrix.shape[0]])
        ax.set_title(label)
        ax.set_xlim(0, audio.duration)

    S = 10 * np.log10(spec_feat(audio) + 1e-9)
    M = 10 * np.log10(mel_feat(audio) + 1e-9)
    MF = mfcc_feat(audio)
    CH = chroma_feat(audio)

    _show(axes[1], S,  "Spectrogram (dB)", cmap="magma")
    _show(axes[2], M,  "Mel Spectrogram (dB)", cmap="magma")
    _show(axes[3], MF, "MFCCs", cmap="coolwarm")
    _show(axes[4], CH, "Chromagram", cmap="viridis")

    for ax in axes:
        ax.set_xlabel("Time (s)")
    plt.suptitle(f"Europa Feature Overview — {audio.path or 'audio'}", y=1.01)
    plt.tight_layout()
    if show:
        plt.show()
    return fig
