"""
renderer.py — Renders WiFi heatmaps using matplotlib.
Supports both floorplan overlay and grid-only mode.

Ekahau-style rendering:
  - Colormap runs red (strongest) → orange → yellow → green → teal → blue →
    violet (weakest), matching professional site survey tools.
  - Alpha channel is derived per-pixel from signal strength so strong-signal
    zones are opaque and weak/absent zones fade to transparent, revealing the
    floorplan or grid beneath. This produces the "bubble" zone appearance.
  - A Gaussian smoothing pass is applied to the RGBA image to soften the
    zone edges and blend overlapping coverage areas naturally.
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import matplotlib.patches as mpatches
import matplotlib.path as mpath
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from matplotlib.patches import Circle
from scipy.ndimage import gaussian_filter
from typing import Optional
from PIL import Image

from data import Session
from interpolator import interpolate, HeatmapData


# ── Colormaps ─────────────────────────────────────────────────────────────────

_EKAHAU_COLORS = [
    (0.45, 0.00, 0.55),   # -90 dBm  violet / purple
    (0.10, 0.10, 0.80),   # -80 dBm  blue
    (0.00, 0.60, 0.80),   # -70 dBm  teal
    (0.10, 0.80, 0.30),   # -65 dBm  green
    (0.70, 0.95, 0.10),   # -58 dBm  yellow-green
    (1.00, 0.85, 0.00),   # -52 dBm  yellow
    (1.00, 0.50, 0.00),   # -46 dBm  orange
    (0.95, 0.10, 0.10),   # -30 dBm  red
]
SIGNAL_CMAP = mcolors.LinearSegmentedColormap.from_list("ekahau", _EKAHAU_COLORS)

VMIN, VMAX = -90, -30


def render_heatmap(
    session: Session,
    bssids: list[str],
    fig: Optional[Figure] = None,
    ax: Optional[Axes] = None,
    alpha: float = 0.85,
    show_points: bool = True,
    show_colorbar: bool = True,
    export_path: Optional[str] = None,
) -> tuple[Figure, Axes, Optional[HeatmapData]]:
    """
    Render an Ekahau-style heatmap for one or more BSSIDs.

    The heatmap is composited as an RGBA image where:
      - RGB  = signal strength mapped through the Ekahau colormap
      - Alpha = derived from signal strength so strong zones are opaque and
                weak/absent zones fade to transparent, revealing the background.

    If any selected BSSIDs have physical positions set in session.ap_positions,
    synthetic anchor points are injected into the interpolation so the heatmap
    peak is correctly located at the real transmitter position.
    """
    if not bssids:
        raise ValueError("At least one BSSID must be provided.")

    if fig is None or ax is None:
        fig, ax = plt.subplots(figsize=(10, 7))
        fig.subplots_adjust(right=0.88)

    # Remove stale colorbar axes from previous renders
    for ax_obj in fig.axes[1:]:
        ax_obj.remove()
    ax.clear()

    points, values          = session.get_points_and_values_multi(bssids)
    missing_points          = session.get_missing_points(bssids)
    anchor_pts, anchor_vals = session.get_anchor_points(bssids)
    w, h                    = session.canvas_width, session.canvas_height

    interp_points = points + anchor_pts
    interp_values = values + anchor_vals

    # ── Title label ───────────────────────────────────────────────────────────
    if len(bssids) == 1:
        ssid     = session.get_ssid(bssids[0])
        ap_label = f"{ssid}  •  {bssids[0]}"
    else:
        ssids    = list({session.get_ssid(b) for b in bssids})
        ssid     = ssids[0] if len(ssids) == 1 else f"{len(ssids)} networks"
        ap_label = f"{ssid}  •  {len(bssids)} BSSIDs averaged"

    # ── Background ────────────────────────────────────────────────────────────
    ax.set_facecolor('#1a1a1a')
    fig.patch.set_facecolor('#0d0d1a')

    if session.floorplan_path:
        try:
            fp_img = Image.open(session.floorplan_path).convert("RGBA")
            fp_img = fp_img.resize((w, h), Image.LANCZOS)
            ax.imshow(np.array(fp_img), extent=[0, w, h, 0],
                      aspect='auto', zorder=1)
        except Exception as e:
            _draw_grid_background(ax, w, h)
            print(f"[renderer] Could not load floorplan: {e}")
    else:
        _draw_grid_background(ax, w, h)

    # ── Heatmap ───────────────────────────────────────────────────────────────
    heatmap_data = None
    MIN_POINTS   = 4

    if len(interp_points) >= MIN_POINTS:
        heatmap_data = interpolate(interp_points, interp_values, w, h)
        grid         = heatmap_data.grid_z
        rgba         = _build_ekahau_rgba(grid, alpha)

        ax.imshow(
            rgba,
            extent=[0, w, 0, h], origin='upper',
            aspect='auto', zorder=2, interpolation='bilinear'
        )

        if show_colorbar:
            _draw_colorbar(fig)

    else:
        collected = len(interp_points)
        remaining = MIN_POINTS - collected
        bar = f"[{'█' * collected}{'░' * remaining}]  {collected}/{MIN_POINTS}"
        msg = (f"Collecting data\n\n{bar}\n\n"
               f"Add {remaining} more point{'s' if remaining != 1 else ''} "
               f"to generate the heatmap.")
        ax.text(w / 2, h / 2, msg,
                ha='center', va='center', fontsize=12,
                color='#8888bb', linespacing=1.9, fontfamily='monospace',
                bbox=dict(boxstyle='round,pad=0.8',
                          facecolor='#10101e', alpha=0.8),
                zorder=3)

    # ── Measurement dots ──────────────────────────────────────────────────────
    radius = max(w, h) * 0.012
    if show_points:
        for (px, py), dbm in zip(points, values):
            norm_val = (dbm - VMIN) / (VMAX - VMIN)
            color    = SIGNAL_CMAP(np.clip(norm_val, 0.0, 1.0))
            ax.add_patch(Circle((px, py), radius=radius,
                                facecolor=color, edgecolor='white',
                                linewidth=1.5, zorder=5, alpha=0.95))
            ax.text(px, py - radius * 1.8, f"{int(round(dbm))}",
                    ha='center', va='bottom', fontsize=7,
                    color='white', fontweight='bold', zorder=6)

        for (px, py) in missing_points:
            ax.add_patch(Circle((px, py), radius=radius,
                                facecolor='black', edgecolor='white',
                                linewidth=1.5, zorder=5, alpha=0.95))
            ax.text(px, py, "?",
                    ha='center', va='center', fontsize=7,
                    color='white', fontweight='bold', zorder=6)

        _draw_ap_markers(ax, session, bssids, radius)

    # ── Labels & axes ─────────────────────────────────────────────────────────
    total      = len(points) + len(missing_points)
    point_note = f"{total} measurement{'s' if total != 1 else ''}"
    if missing_points:
        point_note += f"  •  {len(missing_points)} out of range (●)"
    if len(bssids) > 1:
        point_note += f"  •  avg of {len(bssids)} BSSIDs"
    n_pinned = sum(1 for b in bssids if b in session.ap_positions)
    if n_pinned:
        point_note += f"  •  {n_pinned} AP position{'s' if n_pinned != 1 else ''} pinned (◆)"
    if not session.floorplan_path:
        point_note += "  ⚠ No floorplan — adding one increases accuracy"

    ax.set_title(f"WiFi Heatmap\n{ap_label}",
                 fontsize=12, fontweight='bold', color='white', pad=10)
    ax.set_xlabel(point_note, fontsize=9, color='#888888')
    ax.set_xlim(0, w)
    ax.set_ylim(h, 0)
    ax.set_xticks([])
    ax.set_yticks([])

    if export_path:
        fig.savefig(export_path, dpi=150, bbox_inches='tight',
                    facecolor=fig.get_facecolor())

    return fig, ax, heatmap_data


# ── Colorbar ──────────────────────────────────────────────────────────────────

def _draw_colorbar(fig: Figure):
    """
    Draw the signal strength colorbar by rendering the colormap gradient
    directly as an imshow onto a dedicated axes.

    This approach is backend-agnostic and guaranteed to show color — it does
    not rely on fig.colorbar() interpreting a ScalarMappable, which can silently
    render white on some matplotlib/backend combinations.
    """
    cax = fig.add_axes([0.91, 0.15, 0.025, 0.70])

    # Build the gradient image: 256 rows (dBm steps) × 1 column, RGBA
    gradient = np.linspace(1.0, 0.0, 256).reshape(256, 1)   # top=VMAX, bottom=VMIN
    cax.imshow(gradient, aspect='auto', cmap=SIGNAL_CMAP,
               extent=[0, 1, VMIN, VMAX], origin='upper')

    # Style the axes
    cax.yaxis.set_label_position('right')
    cax.yaxis.tick_right()
    cax.set_ylabel("Signal Strength (dBm)", fontsize=9, color='white', labelpad=8)
    cax.tick_params(axis='y', labelsize=7, colors='white', length=3)
    cax.tick_params(axis='x', which='both', bottom=False, labelbottom=False)
    cax.set_xlim(0, 1)
    cax.set_ylim(VMIN, VMAX)

    # Spine styling
    for spine in cax.spines.values():
        spine.set_edgecolor('#444444')

    # Qualitative labels
    cax.text(1.6, VMAX, "Excellent", transform=cax.get_yaxis_transform(),
             fontsize=7, color='#cccccc', va='top')
    cax.text(1.6, VMIN, "Poor", transform=cax.get_yaxis_transform(),
             fontsize=7, color='#cccccc', va='bottom')

    # "No signal" swatch below the bar
    cax.text(0.5, VMIN - 3, "No signal",
             transform=cax.get_yaxis_transform(),
             fontsize=6, color='white', va='top', ha='center',
             bbox=dict(facecolor='black', edgecolor='#444444',
                       pad=2.0, boxstyle='round'))


# ── RGBA builder ──────────────────────────────────────────────────────────────

def _build_ekahau_rgba(grid: np.ndarray, max_alpha: float) -> np.ndarray:
    """
    Convert a dBm grid into an RGBA uint8 image with Ekahau-style rendering.

    Alpha channel:
      - Signal ≥ VMAX (-30 dBm)  → fully opaque (max_alpha)
      - Signal ≤ FADE_FLOOR      → fully transparent
      - Between                  → cosine ramp for natural bubble-edge blending

    Gaussian blur on alpha blends overlapping AP coverage zones smoothly.
    """
    FADE_FLOOR = -85.0
    BLUR_SIGMA = 8.0

    res_h, res_w = grid.shape

    norm      = np.clip((grid - VMIN) / (VMAX - VMIN), 0.0, 1.0)
    rgb_float = SIGNAL_CMAP(norm)[:, :, :3]

    fade_range = VMAX - FADE_FLOOR
    t          = np.clip((grid - FADE_FLOOR) / fade_range, 0.0, 1.0)
    alpha_raw  = 0.5 * (1.0 - np.cos(np.pi * t))
    alpha_raw[np.isnan(grid)] = 0.0

    alpha_blur = gaussian_filter(alpha_raw, sigma=BLUR_SIGMA)
    alpha_blur = np.clip(alpha_blur * max_alpha, 0.0, max_alpha)

    rgba          = np.zeros((res_h, res_w, 4), dtype=np.uint8)
    rgba[:, :, 0] = (rgb_float[:, :, 0] * 255).astype(np.uint8)
    rgba[:, :, 1] = (rgb_float[:, :, 1] * 255).astype(np.uint8)
    rgba[:, :, 2] = (rgb_float[:, :, 2] * 255).astype(np.uint8)
    rgba[:, :, 3] = (alpha_blur          * 255).astype(np.uint8)

    return rgba


# ── AP position markers ───────────────────────────────────────────────────────

def _draw_ap_markers(ax: Axes, session: Session, bssids: list[str], radius: float):
    """
    Draw a distinct diamond ◆ icon at each AP whose physical position is known.

    Visual design (intentionally different from measurement circles):
      - Outer dashed pulse ring
      - White-bordered diamond
      - Signal-colored inner diamond
      - Three WiFi arc lines above
      - SSID label below
    """
    for bssid in bssids:
        if bssid not in session.ap_positions:
            continue

        px, py = session.ap_positions[bssid]
        ssid   = session.get_ssid(bssid)
        r      = radius * 1.6

        # Pulse ring
        ax.add_patch(Circle((px, py), radius=r * 1.9,
                            facecolor='none', edgecolor='white',
                            linewidth=0.8, zorder=7, alpha=0.35,
                            linestyle='--'))

        # Outer white diamond
        ax.add_patch(mpatches.RegularPolygon(
            (px, py), numVertices=4, radius=r * 1.15,
            orientation=np.pi / 4,
            facecolor='white', edgecolor='white',
            linewidth=0, zorder=8, alpha=0.95))

        # Signal-colored inner diamond
        real_pts, real_vals = session.get_points_and_values(bssid)
        anchor_dbm = min(-25, max(real_vals) + 5) if real_vals else -35
        norm_val   = np.clip((anchor_dbm - VMIN) / (VMAX - VMIN), 0.0, 1.0)
        color      = SIGNAL_CMAP(norm_val)
        ax.add_patch(mpatches.RegularPolygon(
            (px, py), numVertices=4, radius=r * 0.80,
            orientation=np.pi / 4,
            facecolor=color, edgecolor='none',
            zorder=9, alpha=0.95))

        # WiFi arcs
        for arc_r, arc_alpha in [(r * 0.55, 0.9), (r * 0.9, 0.65), (r * 1.25, 0.40)]:
            ax.add_patch(mpatches.Arc(
                (px, py - r * 0.15),
                width=arc_r * 2, height=arc_r * 2,
                angle=0, theta1=30, theta2=150,
                color='white', linewidth=1.4,
                zorder=10, alpha=arc_alpha))

        # SSID label
        ax.text(px, py + r * 2.2, ssid,
                ha='center', va='top', fontsize=7,
                color='white', fontweight='bold', zorder=10,
                bbox=dict(boxstyle='round,pad=0.3',
                          facecolor='#000000', alpha=0.55,
                          edgecolor='none'))

        # Centre glyph
        ax.text(px, py, '◆',
                ha='center', va='center', fontsize=6,
                color='white', zorder=11, alpha=0.7)


# ── Grid background ───────────────────────────────────────────────────────────

def _draw_grid_background(ax: Axes, w: int, h: int):
    """Draw a subtle coordinate grid when no floorplan is loaded."""
    ax.set_facecolor('#1a1a1a')
    spacing = max(w, h) // 10

    for x in range(0, w + 1, spacing):
        ax.axvline(x, color='#2a2a2a', linewidth=0.6, zorder=0)
    for y in range(0, h + 1, spacing):
        ax.axhline(y, color='#2a2a2a', linewidth=0.6, zorder=0)

    for x in range(0, w + 1, spacing * 2):
        ax.text(x, h - 4, str(x), ha='center', va='bottom',
                fontsize=6, color='#3a3a3a')
    for y in range(0, h + 1, spacing * 2):
        ax.text(2, y, str(y), ha='left', va='center',
                fontsize=6, color='#3a3a3a')