"""Independent 3D SAR views of the volume published by the sar3d processor.

Scene axes: X across, Y height above the ground (Y = 0 is the ground), Z along
the ground away from the radar. Payloads share normalization by the maximum
away from the Y boundaries:

- ``sar3d_front``, ``sar3d_side``, ``sar3d_top``: maximum projections onto X-Y
  (looking ahead from the radar), Z-Y (from the side) and X-Z (from above);
- ``sar3d_oblique``: orthographic projection along the camera of the notebook's
  3D scene; ``sar3d_oblique_basis`` (2x3) gives its image axes u and v as scene
  directions, which places the scene box, the ground and the maximum on it;
- ``sar3d_target_width``: axial FWHM [X, Y, Z] in metres through the joint peak.
- ``sar3d_peak``: [level, X, Y, Z] of that maximum before normalization.
- ``sar3d_status``: [code, UTF-8 bytes...] with ready CUDA (0), no signal (1),
  processing error (2), or CUDA unavailable (3), including the explanation.
"""

from __future__ import annotations

import itertools

import numpy as np
import pyqtgraph as pg
from PyQt6.QtCore import QRectF, Qt
from PyQt6.QtWidgets import QGridLayout, QLabel, QVBoxLayout, QWidget

from python_app.gui.controllers.renderers.common import ANTENNA_STYLES, build_lut, create_plot, selected_antenna_coordinates
from python_app.gui.result_payloads import collection_payload_by_name
from python_app.models.dataset_model import ResultCollection, ResultPayload
from python_app.models.gui_profile_schema import GuiSar3dStateModel
from python_app.models.run_config_schema import GprModel

# Status codes published by the sar3d processor in its "sar3d_status" table.
SAR3D_STATES = {0: "ready", 1: "no signal", 2: "error", 3: "warning"}

# payload -> (title, x label, y label, scene axes (0 = X, 1 = Y, 2 = Z) along the plot axes)
SAR3D_VIEWS: dict[str, tuple[str, str, str, tuple[int, int] | None]] = {
    "sar3d_front": ("Front (X-Y)", "X", "Y", (0, 1)),
    "sar3d_side": ("Side (Z-Y)", "Z", "Y", (2, 1)),
    "sar3d_top": ("Top (X-Z)", "X", "Z", (0, 2)),
    "sar3d_oblique": ("3D Projection", "", "", None),
}

# Corners of the scene box indexed by bits (x, y, z); an edge joins corners one bit apart.
_BOX_EDGES = [(a, b) for a, b in itertools.combinations(range(8), 2) if bin(a ^ b).count("1") == 1]
# Axis names go on the edges meeting at the bottom corner nearest the camera (X max, Y min, Z min).
_AXIS_EDGES = {"X": (0, 1), "Y": (1, 3), "Z": (1, 5)}


def sar3d_status(collection: ResultCollection) -> tuple[int, str] | None:
    """Status code (a SAR3D_STATES key) and message, one UTF-8 byte per table column."""
    payload = collection_payload_by_name(collection, "sar3d_status", kind=4)
    if payload is None:
        return None
    if payload.table.ndim != 2 or not payload.table.size or not np.isfinite(payload.table).all():
        return 2, "Invalid SAR3D status table"
    row = payload.table[0]
    code = int(row[0])
    if code not in SAR3D_STATES or code != row[0]:
        return 2, f"Unknown SAR3D status code: {row[0]}"
    return code, bytes(row[1:].astype(np.uint8)).decode("utf-8", errors="replace")


def _image_rect(x_axis: np.ndarray, y_axis: np.ndarray) -> QRectF:
    """Rectangle covering the pixels centred on the axis samples."""
    dx = (x_axis[-1] - x_axis[0]) / max(len(x_axis) - 1, 1)
    dy = (y_axis[-1] - y_axis[0]) / max(len(y_axis) - 1, 1)
    return QRectF(x_axis[0] - dx / 2, y_axis[0] - dy / 2, x_axis[-1] - x_axis[0] + dx, y_axis[-1] - y_axis[0] + dy)


class Sar3dRenderer:
    """Draws the 3D SAR views."""

    def __init__(self) -> None:
        self.widget = QWidget()
        layout = QVBoxLayout(self.widget)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(6)
        self.status_label = QLabel()
        self.status_label.setTextFormat(Qt.TextFormat.PlainText)
        self.status_label.setWordWrap(True)
        layout.addWidget(self.status_label)
        self.grid = QWidget()
        grid_layout = QGridLayout(self.grid)
        grid_layout.setContentsMargins(0, 0, 0, 0)
        grid_layout.setSpacing(6)
        layout.addWidget(self.grid, stretch=1)
        self.plots: dict[str, pg.PlotWidget] = {}
        for index, (name, (title, *_rest)) in enumerate(SAR3D_VIEWS.items()):
            plot = create_plot(title)
            grid_layout.addWidget(plot, index // 2, index % 2)
            self.plots[name] = plot
        lut = build_lut(["#081c15", "#1b4332", "#ffd166", "#f94144"])   # the GPR palette
        lut = np.column_stack([lut, np.full(len(lut), 255, dtype=np.uint8)])
        lut[0, 3] = 0   # below SAR_CLOUD_THRESHOLD
        self.antennas: dict[str, tuple[pg.ScatterPlotItem, pg.ScatterPlotItem]] = {}
        self.images: dict[str, pg.ImageItem] = {}
        self.targets: dict[str, pg.PlotCurveItem] = {}
        for name, (_title, x_label, y_label, _axes) in SAR3D_VIEWS.items():
            plot = self.plots[name]
            plot.setAspectLocked(True)
            if _axes is not None:
                markers = (pg.ScatterPlotItem(), pg.ScatterPlotItem())
                for marker in markers:
                    marker.setZValue(4)
                    plot.addItem(marker)
                self.antennas[name] = markers
            item = plot.getPlotItem()
            assert item is not None
            view = item.getViewBox()
            view.setMouseEnabled(x=False, y=False)
            view.setMenuEnabled(False)
            view.enableAutoRange(x=False, y=False)
            item.hideButtons()
            image = pg.ImageItem(axisOrder="row-major", lut=lut)
            plot.addItem(image)
            # Hollow, so the response at the maximum stays visible.
            peak = pg.PlotCurveItem(pen=pg.mkPen("w", width=1.5))
            peak.setZValue(3)
            plot.addItem(peak)
            self.images[name] = image
            self.targets[name] = peak
            if x_label:
                plot.showGrid(x=True, y=True, alpha=0.2)
                plot.setLabel("bottom", x_label, units="m")
                plot.setLabel("left", y_label, units="m")
            if y_label == "Y":
                ground = pg.InfiniteLine(pos=0.0, angle=0, pen=pg.mkPen("w", width=1, style=Qt.PenStyle.DashLine))
                ground.setZValue(2)
                plot.addItem(ground)

        oblique = self.plots["sar3d_oblique"]
        oblique.hideAxis("bottom")
        oblique.hideAxis("left")
        oblique.setAspectLocked(True)
        # The notebook scene's ground, under the image so it shows through weak voxels.
        self.ground = pg.PlotCurveItem(
            pen=pg.mkPen(115, 105, 88), brush=pg.mkBrush(115, 105, 88, 110), fillLevel="enclosed"
        )
        self.ground.setZValue(-1)
        oblique.addItem(self.ground)
        self.box = pg.PlotCurveItem(pen=pg.mkPen(190, 190, 190, 170), connect="pairs")
        self.box.setZValue(2)
        oblique.addItem(self.box)
        self.axis_labels: dict[str, pg.TextItem] = {}
        for axis_name in _AXIS_EDGES:
            label = pg.TextItem(axis_name, color=(190, 190, 190), anchor=(0.5, 0.5))
            oblique.addItem(label)
            self.axis_labels[axis_name] = label
        self.clear()

    def clear(self, message: str = "SAR3D: waiting for measurements") -> None:
        """Clear all four 3D SAR views."""
        for name, (title, *_rest) in SAR3D_VIEWS.items():
            self.images[name].hide()
            self.targets[name].setData([], [])
            self.plots[name].setTitle(title)
        for markers in self.antennas.values():
            for marker in markers:
                marker.setData([], [])
        self.box.setData([], [])
        self.ground.setData([], [])
        for label in self.axis_labels.values():
            label.hide()
        self.status_label.setText(message)

    def draw(self, collection: ResultCollection, settings: GuiSar3dStateModel, geometry: GprModel) -> bool:
        """Draw four maximum projections of the volume; return success."""
        status = sar3d_status(collection)
        if status is not None and status[0] != 0:
            self.clear(f"SAR3D: {SAR3D_STATES[status[0]]} — {status[1]}")
            return False
        payloads: dict[str, ResultPayload] = {}
        for name in SAR3D_VIEWS:
            payload = collection_payload_by_name(collection, name, kind=3)
            if payload is None or not payload.image.size:
                self.clear("SAR3D: waiting for projection data")
                return False
            payloads[name] = payload
        basis = collection_payload_by_name(collection, "sar3d_oblique_basis", kind=4)
        if basis is None:
            self.clear("SAR3D: waiting for projection data")
            return False
        peak = collection_payload_by_name(collection, "sar3d_peak", kind=4)
        level, *peak_xyz = peak.table[0] if peak is not None else (None,)

        antennas = selected_antenna_coordinates(geometry, settings.input_positions, settings.output_positions)
        pitch = np.radians(settings.look_angle_deg)
        for coordinates in antennas:
            y, z = coordinates[:, 1].copy(), coordinates[:, 2].copy()
            coordinates[:, 1] = settings.bp_ground_depth_m + y*np.cos(pitch) - z*np.sin(pitch)
            coordinates[:, 2] = y*np.sin(pitch) + z*np.cos(pitch)
        width_payload = collection_payload_by_name(collection, "sar3d_target_width", kind=4)
        widths = width_payload.table[0] if width_payload is not None else np.full(3, settings.voxel_m)
        zoom = settings.zoom_percent / 100.0
        for name, payload in payloads.items():
            title, _x_label, _y_label, axes = SAR3D_VIEWS[name]
            rect = _image_rect(payload.image_x_axis, payload.image_y_axis)
            image = self.images[name]
            image.setImage(np.where(payload.image >= settings.cloud_threshold, payload.image, 0),
                           autoLevels=False, levels=(0, 1))
            image.setRect(rect)
            image.show()
            plot = self.plots[name]
            position = (basis.table @ np.asarray(peak_xyz) if axes is None else np.asarray(peak_xyz)[list(axes)]) if peak_xyz else None
            if axes is not None:
                for marker, coordinates, (symbol, brush, pen) in zip(self.antennas[name], antennas, ANTENNA_STYLES):
                    xx, yy = coordinates[:, axes[0]], coordinates[:, axes[1]]
                    marker.setData(x=xx, y=yy, symbol=symbol, size=13, brush=pg.mkBrush(brush), pen=pg.mkPen(pen))
                    if len(xx):
                        pad = settings.voxel_m
                        rect = rect.united(QRectF(xx.min()-pad, yy.min()-pad, np.ptp(xx)+2*pad, np.ptp(yy)+2*pad))
            if position is not None:
                extent = np.abs(basis.table) @ widths if axes is None else widths[list(axes)]
                half = .5 * max(extent)
                self.targets[name].setData(position[0]+np.array([-half, half, half, -half, -half]),
                                           position[1]+np.array([-half, -half, half, half, -half]))
                if zoom > 1.0:
                    center = position + (np.array([rect.center().x(), rect.center().y()])-position) / zoom**2
                    rect = QRectF(center[0]-rect.width()/(2*zoom), center[1]-rect.height()/(2*zoom),
                                  rect.width()/zoom, rect.height()/zoom)
            else:
                self.targets[name].setData([], [])
            plot.setRange(rect, padding=0.04)
            plot.setTitle(title)

        # Scene box and ground on the oblique view, from the extents of the other views.
        front, top = payloads["sar3d_front"], payloads["sar3d_top"]
        bounds = [
            (front.image_x_axis[0], front.image_x_axis[-1]),
            (front.image_y_axis[0], front.image_y_axis[-1]),
            (top.image_y_axis[0], top.image_y_axis[-1]),
        ]
        corners = np.array([[bounds[axis][(index >> axis) & 1] for axis in range(3)] for index in range(8)])
        projected = corners @ basis.table.T
        box = projected[np.ravel(_BOX_EDGES)]
        self.box.setData(box[:, 0], box[:, 1])
        centre = projected.mean(axis=0)
        for axis_name, (start, end) in _AXIS_EDGES.items():
            middle = (projected[start] + projected[end]) / 2
            self.axis_labels[axis_name].setPos(*(middle + 0.08 * (middle - centre)))   # just outside
            self.axis_labels[axis_name].show()
        (x0, x1), _, (z0, z1) = bounds
        ground = (np.array([[x0, 0.0, z0], [x1, 0.0, z0], [x1, 0.0, z1], [x0, 0.0, z1], [x0, 0.0, z0]]) @ basis.table.T).T
        self.ground.setData(ground[0], ground[1])
        summary = f"SAR3D: ready — {status[1]}" if status else "SAR3D: ready"
        if level is not None:
            x, y, z = peak_xyz
            summary += f" | max X {x:+.2f}, Y {y:+.2f}, Z {z:.2f} m | level {level:.3g}"
        self.status_label.setText(f"{summary} | zoom {settings.zoom_percent}%")
        return True