"""Jupyter / Notebook integration.

Provides inline waveform display, signal plotting, and heatmaps
for use in Jupyter notebooks.

Usage:
    t.show(signals=["tb.dut.clk", "tb.dut.counter"], start=0, end=100)
    t.tb.dut.counter.plot(kind="step")
    t.heatmap(signals=[...], start=0, end=200)
"""

from __future__ import annotations
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from ewal.trace import Trace
    from ewal.signal import SignalRef


def show(trace: "Trace", signals: list[str] = None,
         start: int = 0, end: int = None):
    """Render an interactive waveform viewer in a Jupyter cell.

    Generates an HTML/SVG waveform display with:
    - Digital waveform rendering (step function)
    - Bus values displayed as hex
    - Zoom/pan via scroll
    """
    if end is None:
        end = min(trace.max_index, start + 200)

    if signals is None:
        signals = list(trace._core.signal_names())[:10]

    # Build SVG waveform
    sig_height = 30
    label_width = 200
    cycle_width = 8
    total_cycles = end - start
    width = label_width + total_cycles * cycle_width + 20
    height = len(signals) * sig_height + 40

    svg_parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
        f'style="background:white; font-family:monospace; font-size:10px;">',
        '<style>.sig-line{stroke:#333;stroke-width:1.5;fill:none} '
        '.sig-label{fill:#333;font-size:11px} '
        '.bus-val{fill:#666;font-size:8px} '
        '.grid{stroke:#eee;stroke-width:0.5}</style>',
    ]

    for si, sig_name in enumerate(signals):
        y_base = si * sig_height + 30
        w = trace._core.signal_width(sig_name)

        # Label
        short_name = sig_name.split(".")[-1] if "." in sig_name else sig_name
        svg_parts.append(
            f'<text x="5" y="{y_base + 15}" class="sig-label">{short_name}</text>'
        )

        # Waveform
        if w == 1:
            # Digital: draw step waveform
            points = []
            for i in range(total_cycles):
                idx = start + i
                if idx > trace.max_index:
                    break
                val = trace._get_signal_value(sig_name, idx)
                x = label_width + i * cycle_width
                y = y_base + (5 if val else sig_height - 5)
                if points:
                    prev_x, prev_y = points[-1]
                    if y != prev_y:
                        points.append((x, prev_y))  # vertical transition
                points.append((x, y))

            if points:
                path = "M" + " L".join(f"{x},{y}" for x, y in points)
                svg_parts.append(f'<path d="{path}" class="sig-line" />')
        else:
            # Bus: draw boxes with hex values
            prev_val = None
            box_start_x = label_width
            for i in range(total_cycles):
                idx = start + i
                if idx > trace.max_index:
                    break
                val = trace._get_signal_value(sig_name, idx)
                x = label_width + i * cycle_width

                if val != prev_val or i == total_cycles - 1:
                    if prev_val is not None and box_start_x < x:
                        bw = x - box_start_x
                        svg_parts.append(
                            f'<rect x="{box_start_x}" y="{y_base + 2}" '
                            f'width="{bw}" height="{sig_height - 6}" '
                            f'fill="#d4edda" stroke="#333" stroke-width="0.5" rx="2"/>'
                        )
                        if bw > 30:
                            hex_val = f"0x{prev_val:x}" if isinstance(prev_val, int) else str(prev_val)
                            svg_parts.append(
                                f'<text x="{box_start_x + 3}" y="{y_base + sig_height - 8}" '
                                f'class="bus-val">{hex_val}</text>'
                            )
                    box_start_x = x
                    prev_val = val

    svg_parts.append("</svg>")
    svg_html = "\n".join(svg_parts)

    try:
        from IPython.display import display, HTML
        display(HTML(svg_html))
    except ImportError:
        print(svg_html)

    return svg_html


def plot_signal(trace: "Trace", signal_path: str,
                start: int = 0, end: int = None, kind: str = "step"):
    """Plot a signal using matplotlib.

    Args:
        kind: "step" for traditional waveform, "line" for line chart
    """
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib required for plot(). Install with: pip install matplotlib")
        return

    if end is None:
        end = trace.max_index

    indices = list(range(start, min(end + 1, trace.max_index + 1)))
    values = [trace._get_signal_value(signal_path, i) or 0 for i in indices]

    fig, ax = plt.subplots(figsize=(12, 3))

    if kind == "step":
        ax.step(indices, values, where="post", linewidth=1.5)
    else:
        ax.plot(indices, values, linewidth=1.5)

    ax.set_xlabel("Index")
    ax.set_ylabel(signal_path.split(".")[-1])
    ax.set_title(signal_path)
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def heatmap(trace: "Trace", signals: list[str],
            start: int = 0, end: int = None):
    """Display a heatmap of signal values over time."""
    try:
        import matplotlib.pyplot as plt
        import numpy as np
    except ImportError:
        print("matplotlib and numpy required for heatmap()")
        return

    if end is None:
        end = min(trace.max_index, start + 200)

    n_signals = len(signals)
    n_cycles = end - start

    data = np.zeros((n_signals, n_cycles))
    for si, sig in enumerate(signals):
        for ci in range(n_cycles):
            idx = start + ci
            if idx <= trace.max_index:
                val = trace._get_signal_value(sig, idx)
                data[si, ci] = float(val) if val is not None else 0

    fig, ax = plt.subplots(figsize=(14, max(3, n_signals * 0.4)))
    short_names = [s.split(".")[-1] for s in signals]
    im = ax.imshow(data, aspect="auto", cmap="YlOrRd", interpolation="nearest")
    ax.set_yticks(range(n_signals))
    ax.set_yticklabels(short_names)
    ax.set_xlabel("Cycle offset")
    ax.set_title("Signal Heatmap")
    plt.colorbar(im, ax=ax, label="Value")
    plt.tight_layout()
    plt.show()


def _patch_notebook():
    """Add notebook methods to Trace and SignalRef."""
    from ewal.trace import Trace
    from ewal.signal import SignalRef

    def _show(self, signals=None, start=0, end=None):
        return show(self, signals=signals, start=start, end=end)
    Trace.show = _show

    def _heatmap(self, signals, start=0, end=None):
        return heatmap(self, signals=signals, start=start, end=end)
    Trace.heatmap = _heatmap

    def _plot(self, start=0, end=None, kind="step"):
        return plot_signal(self._trace, self._path, start=start, end=end, kind=kind)
    SignalRef.plot = _plot


_patch_notebook()
