"""Core plotting functions for Oscura.

Thin matplotlib wrappers for waveform, spectral, and digital signal visualization.
Each function accepts and returns matplotlib axes for composability.

Example:
    >>> import matplotlib.pyplot as plt
    >>> import oscura as osc
    >>> fig, ax = plt.subplots()
    >>> osc.plot_waveform(ax, trace)
    >>> plt.savefig("waveform.png")
"""

from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, Any

import numpy as np

try:
    import matplotlib.pyplot as plt
    from matplotlib.axes import Axes

    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False


if TYPE_CHECKING:
    from numpy.typing import NDArray

    from oscura.core.types import DigitalTrace, WaveformTrace


def _scale_time(time: NDArray[Any], unit: str) -> tuple[NDArray[Any], str]:
    """Scale time array to appropriate unit."""
    if unit == "auto":
        max_time = time[-1]
        if max_time < 1e-6:
            unit = "ns"
        elif max_time < 1e-3:
            unit = "us"
        elif max_time < 1:
            unit = "ms"
        else:
            unit = "s"
    scale_factors = {"s": 1.0, "ms": 1e3, "us": 1e6, "ns": 1e9}
    return time * scale_factors.get(unit, 1.0), unit


def plot_waveform(
    ax: Axes | None,
    trace: WaveformTrace,
    *,
    time_unit: str = "auto",
    color: str = "#4477AA",
    label: str | None = None,
    show_grid: bool = True,
) -> Axes:
    """Plot time-domain waveform.

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Waveform trace to plot.
        time_unit: Time unit ("s", "ms", "us", "ns", "auto").
        color: Line color.
        label: Legend label.
        show_grid: Show grid lines.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_waveform(ax, trace, time_unit="us")
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 6))

    # Scale time
    time = np.arange(len(trace.data)) / trace.metadata.sample_rate
    time_scaled, time_unit = _scale_time(time, time_unit)

    # Plot
    ax.plot(time_scaled, trace.data, color=color, label=label, linewidth=1.0)
    ax.set_xlabel(f"Time ({time_unit})")
    ax.set_ylabel("Amplitude (V)")
    if show_grid:
        ax.grid(True, alpha=0.3, linestyle="--")
    if label:
        ax.legend()

    return ax


def plot_multi_channel(
    ax: Axes | None,
    traces: list[WaveformTrace],
    *,
    labels: list[str] | None = None,
    colors: list[str] | None = None,
    offset: float | None = None,
    time_unit: str = "auto",
) -> Axes:
    """Plot multiple waveforms on same axes.

    Args:
        ax: Matplotlib axes (creates new if None).
        traces: List of waveform traces.
        labels: Channel labels.
        colors: Line colors.
        offset: Vertical offset between channels (auto if None).
        time_unit: Time unit for x-axis.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_multi_channel(ax, [ch1, ch2, ch3], labels=["CH1", "CH2", "CH3"])
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 8))

    default_colors = ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#AA3377", "#66CCEE"]
    colors = colors or default_colors
    labels = labels or [f"Channel {i + 1}" for i in range(len(traces))]

    # Auto-calculate offset
    if offset is None:
        all_data = np.concatenate([t.data for t in traces])
        offset = (np.max(all_data) - np.min(all_data)) * 1.2

    # Plot each channel
    for i, trace in enumerate(traces):
        time = np.arange(len(trace.data)) / trace.metadata.sample_rate
        time_scaled, time_unit = _scale_time(time, time_unit)
        ax.plot(
            time_scaled,
            trace.data + i * offset,
            color=colors[i % len(colors)],
            label=labels[i] if i < len(labels) else f"Channel {i + 1}",
            linewidth=1.0,
        )

    ax.set_xlabel(f"Time ({time_unit})")
    ax.set_ylabel("Amplitude (V)")
    ax.grid(True, alpha=0.3, linestyle="--")
    ax.legend()
    return ax


def _scale_freq(freq: NDArray[Any], unit: str) -> tuple[NDArray[Any], str]:
    """Scale frequency array to appropriate unit."""
    if unit == "auto":
        max_freq = freq[-1]
        if max_freq > 1e9:
            unit = "GHz"
        elif max_freq > 1e6:
            unit = "MHz"
        elif max_freq > 1e3:
            unit = "kHz"
        else:
            unit = "Hz"
    divisors = {"Hz": 1.0, "kHz": 1e3, "MHz": 1e6, "GHz": 1e9}
    return freq / divisors.get(unit, 1.0), unit


def plot_fft(
    ax: Axes | None,
    trace: WaveformTrace,
    *,
    freq_unit: str = "auto",
    db_scale: bool = True,
    window: str = "hann",
    color: str = "#4477AA",
) -> Axes:
    """Plot FFT magnitude spectrum.

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Waveform trace to analyze.
        freq_unit: Frequency unit ("Hz", "kHz", "MHz", "auto").
        db_scale: Use dB scale for magnitude.
        window: Window function ("hann", "hamming", "blackman", "none").
        color: Line color.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_fft(ax, trace, freq_unit="kHz")
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 6))

    from oscura.analyzers.waveform.spectral import fft

    fft_result = fft(trace, window=window)
    freqs = fft_result[0]
    mags = fft_result[1]
    freqs_scaled, freq_unit = _scale_freq(freqs, freq_unit)

    if db_scale:
        mags = 20 * np.log10(np.abs(mags) + 1e-12)
        ylabel = "Magnitude (dB)"
    else:
        ylabel = "Magnitude"

    ax.plot(freqs_scaled, mags, color=color, linewidth=1.0)
    ax.set_xlabel(f"Frequency ({freq_unit})")
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.3, linestyle="--")
    return ax


def plot_psd(
    ax: Axes | None,
    trace: WaveformTrace,
    *,
    freq_unit: str = "auto",
    window: str = "hann",
    nperseg: int | None = None,
    color: str = "#4477AA",
) -> Axes:
    """Plot power spectral density.

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Waveform trace to analyze.
        freq_unit: Frequency unit ("Hz", "kHz", "MHz", "auto").
        window: Window function.
        nperseg: Segment length for Welch method.
        color: Line color.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_psd(ax, trace)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 6))

    from oscura.analyzers.waveform.spectral import psd

    psd_result = psd(trace, window=window, nperseg=nperseg)
    freqs = psd_result[0]
    power = psd_result[1]
    freqs_scaled, freq_unit = _scale_freq(freqs, freq_unit)

    ax.plot(freqs_scaled, 10 * np.log10(power + 1e-12), color=color, linewidth=1.0)
    ax.set_xlabel(f"Frequency ({freq_unit})")
    ax.set_ylabel("Power Spectral Density (dB/Hz)")
    ax.grid(True, alpha=0.3, linestyle="--")
    return ax


def plot_spectrogram(
    ax: Axes | None,
    trace: WaveformTrace,
    *,
    nfft: int = 1024,
    noverlap: int | None = None,
    cmap: str = "viridis",
) -> Axes:
    """Plot spectrogram (time-frequency representation).

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Waveform trace to analyze.
        nfft: FFT window size.
        noverlap: Overlap samples (default: nfft // 2).
        cmap: Colormap name.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_spectrogram(ax, trace, nfft=512)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 6))

    if noverlap is None:
        noverlap = nfft // 2

    # Generate spectrogram
    sample_rate = trace.metadata.sample_rate
    ax.specgram(trace.data, Fs=sample_rate, NFFT=nfft, noverlap=noverlap, cmap=cmap)

    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Frequency (Hz)")

    return ax


def plot_eye(
    ax: Axes | None,
    trace: WaveformTrace,
    *,
    symbol_rate: float,
    n_symbols: int = 100,
    color: str = "#4477AA",
    alpha: float = 0.1,
) -> Axes:
    """Plot eye diagram for digital signals.

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Waveform trace.
        symbol_rate: Symbol rate in Hz.
        n_symbols: Number of symbols to overlay.
        color: Line color.
        alpha: Line transparency.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_eye(ax, trace, symbol_rate=1e6)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(8, 6))

    # Calculate samples per symbol
    sample_rate = trace.metadata.sample_rate
    samples_per_symbol = int(sample_rate / symbol_rate)

    # Overlay multiple symbols
    for i in range(n_symbols):
        start = i * samples_per_symbol
        end = start + 2 * samples_per_symbol  # Two symbols for eye
        if end > len(trace.data):
            break

        segment = trace.data[start:end]
        time = np.arange(len(segment)) / sample_rate * 1e6  # Convert to microseconds
        ax.plot(time, segment, color=color, alpha=alpha, linewidth=0.8)

    ax.set_xlabel("Time (us)")
    ax.set_ylabel("Amplitude (V)")
    ax.set_title("Eye Diagram")
    ax.grid(True, alpha=0.3)

    return ax


def plot_timing(
    ax: Axes | None,
    trace: DigitalTrace,
    *,
    time_unit: str = "us",
    color: str = "#4477AA",
    label: str | None = None,
) -> Axes:
    """Plot digital timing diagram.

    Args:
        ax: Matplotlib axes (creates new if None).
        trace: Digital trace.
        time_unit: Time unit ("s", "ms", "us", "ns").
        color: Line color.
        label: Signal label.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_timing(ax, digital_trace)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(10, 4))

    time = np.arange(len(trace.data)) / trace.metadata.sample_rate
    time_scaled, time_unit = _scale_time(time, time_unit)

    ax.step(
        time_scaled, trace.data.astype(float), where="post", color=color, label=label, linewidth=1.5
    )
    ax.fill_between(time_scaled, 0, trace.data.astype(float), step="post", alpha=0.2, color=color)

    ax.set_xlabel(f"Time ({time_unit})")
    ax.set_ylabel("Logic Level")
    ax.set_ylim(-0.2, 1.3)
    ax.set_yticks([0, 1])
    ax.set_yticklabels(["LOW", "HIGH"])
    ax.grid(True, alpha=0.3, axis="x")
    if label:
        ax.legend()
    return ax


def plot_histogram(
    ax: Axes | None,
    data: NDArray[Any],
    *,
    bins: int = 50,
    color: str = "#4477AA",
    show_stats: bool = True,
) -> Axes:
    """Plot amplitude histogram with optional statistics.

    Args:
        ax: Matplotlib axes (creates new if None).
        data: Signal data array.
        bins: Number of histogram bins.
        color: Bar color.
        show_stats: Show mean/std overlays.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_histogram(ax, trace.data, bins=100)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(8, 6))

    ax.hist(data, bins=bins, color=color, alpha=0.7, edgecolor="black", linewidth=0.5)

    if show_stats:
        mean_val, std_val = np.mean(data), np.std(data)
        ax.axvline(
            mean_val, color="#EE6677", linestyle="--", linewidth=2, label=f"Mean: {mean_val:.3f}"
        )
        ax.axvline(
            mean_val + std_val,
            color="#AAA",
            linestyle=":",
            linewidth=1.5,
            label=f"+/-1std: {std_val:.3f}",
        )
        ax.axvline(mean_val - std_val, color="#AAA", linestyle=":", linewidth=1.5)
        ax.legend()

    ax.set_xlabel("Amplitude (V)")
    ax.set_ylabel("Count")
    ax.grid(True, alpha=0.3, axis="y")
    return ax


def plot_xy(
    ax: Axes | None,
    trace_x: WaveformTrace,
    trace_y: WaveformTrace,
    *,
    color: str = "#4477AA",
    marker: str = "o",
    markersize: float = 3.0,
) -> Axes:
    """Plot XY scatter (Lissajous pattern).

    Args:
        ax: Matplotlib axes (creates new if None).
        trace_x: X-axis trace.
        trace_y: Y-axis trace.
        color: Marker color.
        marker: Marker style.
        markersize: Marker size.

    Returns:
        Matplotlib axes object.

    Example:
        >>> fig, ax = plt.subplots()
        >>> plot_xy(ax, ch1, ch2)
    """
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")

    if ax is None:
        _, ax = plt.subplots(figsize=(8, 8))

    min_len = min(len(trace_x.data), len(trace_y.data))
    ax.plot(
        trace_x.data[:min_len],
        trace_y.data[:min_len],
        marker=marker,
        linestyle="",
        color=color,
        markersize=markersize,
        alpha=0.5,
    )

    ax.set_xlabel("X Amplitude (V)")
    ax.set_ylabel("Y Amplitude (V)")
    ax.set_aspect("equal")
    ax.grid(True, alpha=0.3)
    return ax


def _removed_plot_stub(name: str) -> None:
    """Emit deprecation warning and raise for removed plot functions.

    Args:
        name: Name of the removed function.

    Raises:
        NotImplementedError: Always, after emitting DeprecationWarning.
    """
    warnings.warn(
        f"{name}() was removed in the visualization refactor. "
        f"Use plot_waveform, plot_fft, plot_eye, plot_timing, or plot_histogram instead.",
        DeprecationWarning,
        stacklevel=3,
    )
    raise NotImplementedError(
        f"{name}() was removed in the visualization refactor. "
        f"Available alternatives: plot_waveform, plot_fft, plot_psd, plot_spectrogram, "
        f"plot_eye, plot_timing, plot_histogram, plot_xy, plot_multi_channel."
    )


# Stub functions for backwards compatibility with demos
def plot_bathtub(*args: Any, **kwargs: Any) -> Any:
    """Removed in visualization refactor. Use plot_eye() for eye diagram analysis.

    .. deprecated::
        This function was removed in the visualization refactor.
        Use :func:`plot_eye` for eye diagram / bathtub curve analysis.

    Raises:
        NotImplementedError: Always raised with migration guidance.
    """
    _removed_plot_stub("plot_bathtub")


def plot_protocol_decode(*args: Any, **kwargs: Any) -> Any:
    """Removed in visualization refactor. Use plot_timing() for protocol visualization.

    .. deprecated::
        This function was removed in the visualization refactor.
        Use :func:`plot_timing` for digital protocol visualization.

    Raises:
        NotImplementedError: Always raised with migration guidance.
    """
    _removed_plot_stub("plot_protocol_decode")


def plot_state_machine(*args: Any, **kwargs: Any) -> Any:
    """Removed in visualization refactor. No direct replacement available.

    .. deprecated::
        This function was removed in the visualization refactor.
        For state machine visualization, use graphviz or networkx directly.

    Raises:
        NotImplementedError: Always raised with migration guidance.
    """
    _removed_plot_stub("plot_state_machine")


def plot_thd_bars(*args: Any, **kwargs: Any) -> Any:
    """Removed in visualization refactor. Use plot_fft() for harmonic analysis.

    .. deprecated::
        This function was removed in the visualization refactor.
        Use :func:`plot_fft` for frequency-domain harmonic analysis.

    Raises:
        NotImplementedError: Always raised with migration guidance.
    """
    _removed_plot_stub("plot_thd_bars")


def plot_quality_summary(*args: Any, **kwargs: Any) -> Any:
    """Removed in visualization refactor. Use the reporting module instead.

    .. deprecated::
        This function was removed in the visualization refactor.
        Use :func:`oscura.reporting.analyze.analyze` for comprehensive quality reports.

    Raises:
        NotImplementedError: Always raised with migration guidance.
    """
    _removed_plot_stub("plot_quality_summary")


def plot_logic_analyzer(*args: Any, **kwargs: Any) -> Axes:
    """Alias for plot_timing for backwards compatibility."""
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")
    # Redirect to plot_timing
    return plot_timing(*args, **kwargs)


def plot_spectrum(*args: Any, **kwargs: Any) -> Axes:
    """Alias for plot_fft for backwards compatibility."""
    if not HAS_MATPLOTLIB:
        raise ImportError("matplotlib required for plotting")
    return plot_fft(*args, **kwargs)


__all__ = [
    "plot_bathtub",
    "plot_eye",
    "plot_fft",
    "plot_histogram",
    "plot_logic_analyzer",
    "plot_multi_channel",
    "plot_protocol_decode",
    "plot_psd",
    "plot_quality_summary",
    "plot_spectrogram",
    "plot_spectrum",
    "plot_state_machine",
    "plot_thd_bars",
    "plot_timing",
    "plot_waveform",
    "plot_xy",
]
