"""QQuickPaintedItem wafer map — ported from the replay app's ReplayWidget.

Draws:
  • Radial ring template (concentric guides + crosshair axes + top notch)
  • RBF heatmap layer (blended under markers via `blend` 0→1)
  • Sensor marker circles colored by band (low/in_range/high)
  • Numbered labels (toggle via `showLabels`)

All sensor coordinates are center-origin mm (from wafer_layouts or a loaded CSV).
"""
from __future__ import annotations

import logging
import math

import numpy as np
from PySide6.QtCore import Property, QPoint, Qt, Signal, Slot
from PySide6.QtGui import (
    QBrush,
    QColor,
    QFont,
    QImage,
    QPainter,
    QPen,
    QPolygon,
)  # fmt: skip
from PySide6.QtQml import QmlElement
from PySide6.QtQuick import QQuickPaintedItem

from pygui.backend.visualization.rbf_heatmap import interpolate_field
from pygui.backend.wafer.zwafer_models import Sensor

log = logging.getLogger(__name__)

QML_IMPORT_NAME = "ISC.Wafer"
QML_IMPORT_MAJOR_VERSION = 1


@QmlElement
class WaferMapItem(QQuickPaintedItem):
    """Painted wafer map; driven by SessionController via QML property bindings."""

    sensorsChanged  = Signal()
    valuesChanged   = Signal()
    bandsChanged    = Signal()
    targetChanged   = Signal()
    marginChanged   = Signal()
    blendChanged    = Signal()
    showLabelsChanged = Signal()
    colorsChanged   = Signal()
    shapeChanged    = Signal()
    sizeChanged     = Signal()
    thicknessChanged = Signal()
    showThicknessChanged = Signal()


    def __init__(self, parent=None):
        super().__init__(parent)
        self._sensors:  list[Sensor] = []
        self._values:   list[float]  = []
        self._bands:    list[str]    = []
        self._target:   float = 149.0
        self._margin:   float = 1.0
        self._blend:    float = 0.0
        self._show_labels: bool = True
        self._shape: str = "round"
        self._size: float = 300.0
        self._thickness_data: list[float] = []
        self._show_thickness: bool = False
        self._thickness_heatmap:QImage | None = None

        # Dark-theme color defaults (match Theme.qml tokens)
        self._ring_color     = QColor("#2A3441")   # waferRingColor  (toneBorder)
        self._axis_color     = QColor("#3A4D5C")   # waferAxisColor  (softBorder)
        self._low_color      = QColor("#5B9DF5")   # sensorLow
        self._in_range_color = QColor("#22C55E")   # sensorInRange
        self._high_color     = QColor("#EF4444")   # sensorHigh
        self._text_color     = QColor("#CBD5E1")   # bodyColor

        # Internal draw state
        self._markers: dict[int, tuple[int, int]] = {}  # sensor index → (px, py)
        self._marker_r: int = 4
        self._heatmap:  QImage | None = None

        self.widthChanged.connect(self._on_resize)
        self.heightChanged.connect(self._on_resize)

    # ── Qt properties ────────────────────────────────────────────────────

    @Property("QVariantList", notify=sensorsChanged)
    def sensors(self) -> list:
        return [{"label": s.label, "x": s.x, "y": s.y} for s in self._sensors]

    @sensors.setter
    def sensors(self, val: list) -> None:
        self._sensors = [
            Sensor(
                label=d["label"],
                x=float(d["x"]),
                y=float(d["y"]),
                side=d.get("side", "right"),
                offset_x=float(d.get("offset_x", 0.0)),
                offset_y=float(d.get("offset_y", 0.0))
            )
            for d in (val or [])
        ]
        self._rebuild()
        self.sensorsChanged.emit()

    @Property("QVariantList", notify=valuesChanged)
    def values(self) -> list:
        return self._values

    @values.setter
    def values(self, val: list) -> None:
        self._values = list(val or [])
        self._rebuild_heatmap()
        self.valuesChanged.emit()
        self.update()

    @Property("QVariantList", notify=bandsChanged)
    def bands(self) -> list:
        return self._bands

    @bands.setter
    def bands(self, val: list) -> None:
        self._bands = list(val or [])
        self.bandsChanged.emit()
        self.update()

    @Property(float, notify=targetChanged)
    def target(self) -> float:
        return self._target

    @target.setter
    def target(self, val: float) -> None:
        self._target = float(val)
        self._rebuild_heatmap()
        self.targetChanged.emit()
        self.update()

    @Property(float, notify=marginChanged)
    def margin(self) -> float:
        return self._margin

    @margin.setter
    def margin(self, val: float) -> None:
        self._margin = float(val)
        self._rebuild_heatmap()
        self.marginChanged.emit()
        self.update()

    @Property(float, notify=blendChanged)
    def blend(self) -> float:
        return self._blend

    @blend.setter
    def blend(self, val: float) -> None:
        self._blend = max(0.0, min(1.0, float(val)))
        if self._blend > 0 and self._heatmap is None:
            self._rebuild_heatmap()
        self.blendChanged.emit()
        self.update()

    @Property(bool, notify=showLabelsChanged)
    def showLabels(self) -> bool:
        return self._show_labels

    @showLabels.setter
    def showLabels(self, val: bool) -> None:
        self._show_labels = bool(val)
        self.showLabelsChanged.emit()
        self.update()

    @Property(str, notify=shapeChanged)
    def shape(self) -> str:
        return self._shape

    @shape.setter
    def shape(self, val: str) -> None:
        self._shape = str(val).lower()
        self._rebuild()
        self.shapeChanged.emit()

    @Property(float, notify=sizeChanged)
    def size(self) -> float:
        return self._size

    @size.setter
    def size(self, val: float) -> None:
        self._size = float(val)
        self._rebuild()
        self.sizeChanged.emit()

    @Property("QVariantList", notify=thicknessChanged)
    def thicknessData(self) -> list:
        return self._thickness_data

    @thicknessData.setter
    def thicknessData(self, val:list) -> None:
        self._thickness_data = list(val or [])
        self._rebuild_thickness()
        self.thicknessChanged.emit()
        self.update

    @Property(bool, notify=showThicknessChanged)
    def showThickness(self) -> bool:
        return self._show_thickness

    @showThickness.setter
    def showThickness(self, val: bool) -> None:
        self._show_thickness = bool(val)
        self.showThicknessChanged.emit()
        self.update()

    # Colour properties — QML can bind these to Theme tokens
    @Property(QColor, notify=colorsChanged)
    def ringColor(self) -> QColor: return self._ring_color
    @ringColor.setter
    def ringColor(self, c: QColor) -> None: self._ring_color = c; self.update()

    @Property(QColor, notify=colorsChanged)
    def axisColor(self) -> QColor: return self._axis_color
    @axisColor.setter
    def axisColor(self, c: QColor) -> None: self._axis_color = c; self.update()

    @Property(QColor, notify=colorsChanged)
    def lowColor(self) -> QColor: return self._low_color
    @lowColor.setter
    def lowColor(self, c: QColor) -> None: self._low_color = c; self.update()

    @Property(QColor, notify=colorsChanged)
    def inRangeColor(self) -> QColor: return self._in_range_color
    @inRangeColor.setter
    def inRangeColor(self, c: QColor) -> None: self._in_range_color = c; self.update()

    @Property(QColor, notify=colorsChanged)
    def highColor(self) -> QColor: return self._high_color
    @highColor.setter
    def highColor(self, c: QColor) -> None: self._high_color = c; self.update()

    @Property(QColor, notify=colorsChanged)
    def textColor(self) -> QColor: return self._text_color
    @textColor.setter
    def textColor(self, c: QColor) -> None: self._text_color = c; self.update()

    # ── slots ─────────────────────────────────────────────────────────────

    @Slot(float, float, result=int)
    def which_marker(self, x: float, y: float) -> int:
        """Return the sensor index nearest to (x, y) within marker radius, else -1."""
        r = max(self._marker_r, 6)
        for idx, (mx, my) in self._markers.items():
            if abs(mx - x) <= r and abs(my - y) <= r:
                return idx
        return -1

    @Slot(str, result=bool)
    def export_image(self, file_path: str) -> bool:
        """Export the current wafer map rendering to a PNG file

        Args:
            file_pat: Absolute path for the output PNG.

        Returns:
        True on success, False on failure
        """
        if not file_path:
            return False
        try:
            result = self.grabToImage()
            img = result.image()
            img.save(file_path, "PNG")
            return True
        except Exception as e:
            log.error("Export failed: %s", e)
            return False

    # ── internal ─────────────────────────────────────────────────────────

    def _rebuild_thickness(self) -> None:
        """Interpolate thickness data into a gray/orange heatmap QImage."""
        if not self._sensors or not self._thickness_data:
            self._thickness_heatmap = None
            return
        ds = self._draw_size()
        r_mm = self._wafer_radius_mm()
        xs = np.array([s.x for s in self._sensors])
        ys = np.array([s.y for s in self._sensors])
        vs = np.array(self._thickness_data[:len(self._sensors)], dtype=float)
        if len(vs) < len(self._sensors):
            self._thickness_heatmap = None
            return
        try:
            field = interpolate_field(
                xs, ys, vs,
                width=ds, height=ds,
                extent=(-r_mm, r_mm, -r_mm, r_mm),
                round_clip=(self._shape == "round"),
            )
        except Exception:
            self._thickness_heatmap = None
            return
        # Gray/orange colormap: map field range 0→1 to gray→orange
        vmin, vmax = np.nanmin(field), np.nanmax(field)
        span = vmax - vmin or 1.0
        t = np.clip((field - vmin) / span, 0.0, 1.0)

        # Gray (0.5, 0.5, 0.5) → orange (1.0, 0.65, 0.0)
        rgb = np.zeros((ds, ds, 3), dtype=np.float32)
        rgb[:, :, 0] = 0.5 + 0.5 * t
        rgb[:, :, 1] = 0.5 + 0.15 * t
        rgb[:, :, 2] = 0.5 - 0.5 * t

        rgb = np.nan_to_num(rgb, nan=0.0)
        rgba = np.zeros((ds, ds, 4), dtype=np.uint8)
        rgba[:, :, :3] = (rgb * 255).astype(np.uint8)
        rgba[:, :, 3] = np.where(np.isfinite(field), 180, 0).astype(np.uint8)

        self._thickness_heatmap = (
            QImage(rgba.tobytes(), ds, ds, QImage.Format.Format_RGBA8888).copy()
        )

    def _on_resize(self) -> None:
        self._rebuild()

    def _rebuild(self) -> None:
        self._compute_markers()
        self._rebuild_heatmap()
        self.update()

    def _draw_size(self) -> int:
        return max(1, int(min(self.width(), self.height())))

    def _center(self) -> tuple[int, int]:
        return int(self.width() / 2), int(self.height() / 2)

    def _wafer_radius_mm(self) -> float:
        """Radius of the wafer bounding circle in mm (5% padding beyond outermost sensor)."""
        if not self._sensors:
            return 150.0
        r = max(math.hypot(s.x, s.y) for s in self._sensors)
        return r * 1.05

    def _sensor_ring_radii_mm(self) -> list[float]:
        """Distinct radial distances of sensor groups, sorted ascending, plus the outer boundary."""
        if not self._sensors:
            r = self._wafer_radius_mm()
            return [r * f for f in (0.25, 0.50, 0.75, 1.0)]
        # Cluster radii that are within 2 mm of each other into one ring; skip center point.
        radii = sorted(r for r in {math.hypot(s.x, s.y) for s in self._sensors} if r > 1.0)
        groups: list[float] = []
        for r in radii:
            if not groups or r - groups[-1] > 2.0:
                groups.append(r)
            else:
                groups[-1] = (groups[-1] + r) / 2  # merge close values
        # Always include the outer boundary ring so the wafer circle is drawn.
        outer = self._wafer_radius_mm()
        if not groups or outer - groups[-1] > 2.0:
            groups.append(outer)
        return groups

    def _scale(self, ds: int, r_mm: float) -> float:
        """Pixels per mm. The wafer radius maps to ds//2 - 24 px."""
        return (ds / 2 - 24) / r_mm

    def _to_px(self, x_mm: float, y_mm: float, cx: int, cy: int, scale: float) -> tuple[int, int]:
        """Center-origin mm → pixel (top-left origin). Y is flipped."""
        return cx + int(x_mm * scale), cy - int(y_mm * scale)

    def _compute_markers(self) -> None:
        ds = self._draw_size()
        r_mm = self._wafer_radius_mm()
        sc = self._scale(ds, r_mm)
        cx, cy = self._center()
        self._marker_r = max(3, ds // 70)
        self._markers = {i: self._to_px(s.x, s.y, cx, cy, sc)
                         for i, s in enumerate(self._sensors)}

    def _rebuild_heatmap(self) -> None:
        if not self._sensors or not self._values or self._blend == 0.0:
            self._heatmap = None
            return
        ds = self._draw_size()
        r_mm = self._wafer_radius_mm()
        xs = np.array([s.x for s in self._sensors])
        ys = np.array([s.y for s in self._sensors])
        vs = np.array(self._values[:len(self._sensors)], dtype=float)
        if len(vs) < len(self._sensors):
            self._heatmap = None
            return
        try:
            field = interpolate_field(
                xs, ys, vs,
                width=ds, height=ds,
                extent=(-r_mm, r_mm, -r_mm, r_mm),
                round_clip=(self._shape == "round"),
            )
        except Exception:
            self._heatmap = None
            return
        self._heatmap = self._field_to_qimage(field, ds)

    def _field_to_qimage(self, field: np.ndarray, ds: int) -> QImage:
        """Apply a band-aware tri-color gradient → RGBA QImage."""
        lo_b = self._target - self._margin
        hi_b = self._target + self._margin
        span = hi_b - lo_b or 1.0
        # t: 0 = lo_b, 1 = hi_b (clipped)
        t = np.clip((field - lo_b) / span, 0.0, 1.0)

        def c(q: QColor) -> np.ndarray:
            return np.array([q.redF(), q.greenF(), q.blueF()], dtype=np.float32)

        lo_c = c(self._low_color)
        mid_c = c(self._in_range_color)
        hi_c = c(self._high_color)

        t2 = t * 2  # 0→2 across full range

        lower = t <= 0.5
        t_lo = np.clip(t2,       0.0, 1.0)[:, :, np.newaxis]  # 0→1 in lower half
        t_hi = np.clip(t2 - 1.0, 0.0, 1.0)[:, :, np.newaxis]  # 0→1 in upper half

        rgb = np.where(
            lower[:, :, np.newaxis],
            lo_c * (1 - t_lo) + mid_c * t_lo,
            mid_c * (1 - t_hi) + hi_c * t_hi,
        )
        # Outside the wafer circle `field` is NaN → NaN propagates into rgb. Alpha
        # masks those pixels anyway, but zero them so the uint8 cast is well-defined.
        rgb = np.nan_to_num(rgb, nan=0.0)
        rgba = np.zeros((ds, ds, 4), dtype=np.uint8)
        rgba[:, :, :3] = (rgb * 255).astype(np.uint8)
        rgba[:, :, 3]  = np.where(np.isfinite(field), 210, 0).astype(np.uint8)

        return QImage(rgba.tobytes(), ds, ds, QImage.Format.Format_RGBA8888).copy()

    # ── paint ─────────────────────────────────────────────────────────────

    def paint(self, painter: QPainter) -> None:
        ds   = self._draw_size()
        r_px = int(ds / 2 - 4)
        cx, cy = self._center()

        painter.setRenderHint(QPainter.RenderHint.Antialiasing)

        self._paint_template(painter, cx, cy, r_px)

        if self._heatmap and self._blend > 0.0:
            painter.setOpacity(self._blend)
            painter.drawImage(cx - self._heatmap.width() // 2,
                              cy - self._heatmap.height() // 2, self._heatmap)
            painter.setOpacity(1.0)

        if self._show_thickness and self._thickness_heatmap:
            painter.setOpacity(0.4)
            painter.drawImage(cx - self._thickness_heatmap.width() // 2,
            cy - self._thickness_heatmap.height()//2, self._thickness_heatmap)
            painter.setOpacity(1.0)

        self._paint_markers(painter)

    def _paint_template(self, painter: QPainter, cx: int, cy: int, r_px: int) -> None:
        ds = self._draw_size()
        r_mm = self._wafer_radius_mm()
        sc = self._scale(ds, r_mm)

        if self._shape == "square":
            # Draw square boundary (thick pen)
            border_pen = QPen(self._ring_color, 2, Qt.PenStyle.SolidLine)
            painter.setPen(border_pen)
            half_size_px = int(self._size / 2 * sc)
            painter.drawRect(cx - half_size_px, cy - half_size_px, 2 * half_size_px, 2 * half_size_px)

            # Crosshair axes
            axis_pen = QPen(self._axis_color, 1, Qt.PenStyle.DashLine)
            painter.setPen(axis_pen)
            painter.drawLine(cx, cy - half_size_px, cx, cy + half_size_px)
            painter.drawLine(cx - half_size_px, cy, cx + half_size_px, cy)

            # Draw concentric square guide lines at the distinct radii of sensor rings
            grid_pen = QPen(self._ring_color, 1, Qt.PenStyle.SolidLine)
            painter.setPen(grid_pen)
            for ring_r_mm in self._sensor_ring_radii_mm()[:-1]:  # exclude the outermost border
                rr = max(1, int(ring_r_mm * sc))
                painter.drawRect(cx - rr, cy - rr, 2 * rr, 2 * rr)
        else:
            # Concentric rings
            ring_pen = QPen(self._ring_color, 1, Qt.PenStyle.SolidLine)
            painter.setPen(ring_pen)
            for ring_r_mm in self._sensor_ring_radii_mm():
                rr = max(1, int(ring_r_mm * sc))
                painter.drawEllipse(cx - rr, cy - rr, 2 * rr, 2 * rr)

            # Crosshair axes
            axis_pen = QPen(self._axis_color, 1, Qt.PenStyle.DashLine)
            painter.setPen(axis_pen)
            painter.drawLine(cx, cy - r_px, cx, cy + r_px)
            painter.drawLine(cx - r_px, cy, cx + r_px, cy)

            # Top notch triangle (wafer orientation marker)
            nw = max(6, ds // 25)
            nh = max(4, ds // 35)
            notch = QPolygon([
                QPoint(cx,           cy - r_px),
                QPoint(cx - nw // 2, cy - r_px + nh),
                QPoint(cx + nw // 2, cy - r_px + nh),
            ])
            painter.setPen(Qt.PenStyle.NoPen)
            painter.setBrush(QBrush(self._axis_color))
            painter.drawPolygon(notch)

    def _paint_markers(self, painter: QPainter) -> None:
        r = self._marker_r

        # Scale font size based on the number of sensors to prevent overlap on dense wafers
        num_sensors = len(self._sensors)
        font_scale = 1.0
        if num_sensors > 60:
            font_scale = 0.7  # reduce font size by 30% for dense wafers
        elif num_sensors > 40:
            font_scale = 0.85

        id_font = QFont()
        id_font.setPointSize(max(4, int(r * font_scale)))
        id_font.setBold(True)

        temp_font = QFont()
        temp_font.setPointSize(max(3, int((r - 1) * font_scale)))

        band_color = {
            "in_range": self._in_range_color,
            "high":     self._high_color,
            "low":      self._low_color,
        }

        for i, s in enumerate(self._sensors):
            if i not in self._markers:
                continue
            px, py = self._markers[i]
            color = band_color.get(
                self._bands[i] if i < len(self._bands) else "in_range",
                self._in_range_color,
            )
            # Filled circle with thin dark outline for contrast over heatmap
            painter.setPen(QPen(QColor(0, 0, 0, 100), 1))
            painter.setBrush(QBrush(color))
            painter.drawEllipse(px - r, py - r, 2 * r, 2 * r)

            if self._show_labels:
                has_temp = i < len(self._values)

                # Fetch text alignment side and offsets
                side = getattr(s, "side", "right").lower()
                ox = getattr(s, "offset_x", 0.0) * r
                oy = getattr(s, "offset_y", 0.0) * r

                # Pre-compute metrics using current scaled fonts
                painter.setFont(id_font)
                id_fm = painter.fontMetrics()
                id_line_h = id_fm.height()
                id_ascent = id_fm.ascent()

                painter.setFont(temp_font)
                temp_fm = painter.fontMetrics()
                temp_line_h = temp_fm.height()
                temp_ascent = temp_fm.ascent()

                id_text = s.label
                temp_text = f"{self._values[i]:.2f}" if has_temp else ""

                id_w = id_fm.horizontalAdvance(id_text)
                temp_w = temp_fm.horizontalAdvance(temp_text) if has_temp else 0
                text_w = max(id_w, temp_w)

                # Height of the 1 or 2-line text block
                text_h = (id_line_h + temp_line_h) if has_temp else id_line_h

                # Calculate box top-left (lx, ly) relative to dot center (px, py)
                gap = 3
                if side == "left":
                    lx = px - r - gap - text_w - ox
                    ly = py - text_h // 2 + oy
                elif side == "top":
                    lx = px - text_w // 2 + ox
                    ly = py - r - gap - text_h - oy
                elif side == "bottom":
                    lx = px - text_w // 2 + ox
                    ly = py + r + gap + oy
                else: # "right" or default
                    lx = px + r + gap + ox
                    ly = py - text_h // 2 + oy

                # Draw Sensor ID (first line)
                painter.setFont(id_font)
                painter.setPen(QPen(self._text_color))
                y1 = ly + id_ascent
                if side in ("top", "bottom"):
                    painter.drawText(lx + (text_w - id_w) // 2, y1, id_text)
                elif side == "left":
                    painter.drawText(lx + (text_w - id_w), y1, id_text)
                else:
                    painter.drawText(lx, y1, id_text)

                # Draw Temperature (second line)
                if has_temp:
                    painter.setFont(temp_font)
                    painter.setPen(QPen(color))
                    y2 = ly + id_line_h + temp_ascent
                    if side in ("top", "bottom"):
                        painter.drawText(lx + (text_w - temp_w) // 2, y2, temp_text)
                    elif side == "left":
                        painter.drawText(lx + (text_w - temp_w), y2, temp_text)
                    else:
                        painter.drawText(lx, y2, temp_text)