"""Predicted-vs-GT height-over-time line plot.

For the per-voxel volume model: the model no longer regresses height
separately (no marginal-Z panel), but height still falls out of the
joint argmax over the volume's ``(Z, P, P)`` cells. We log a matplotlib
line plot of:

  x — predicted timestep ``t ∈ [0, n_window)``
  y — height (m) from the bin center of the joint argmax cell

Two lines: **GT** (from the dataset's target_z) and **Pred** (from the
predicted volume argmax). Y-axis is locked to the dataset's ``z_range``
``(z_lo, z_hi)`` so the chart's scale is comparable across logged steps.
"""
from __future__ import annotations

import cv2
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


def height_curve_panel(
    pred_z_bins: np.ndarray,      # (T,) int  — predicted Z bin per timestep
    gt_z_bins: np.ndarray,        # (T,) int  — GT Z bin per timestep
    z_lo: float,
    z_hi: float,
    n_z: int,
    out_wh: tuple[int, int] = (512, 256),
) -> np.ndarray:
    """Return BGR uint8 panel of pred vs GT height over time."""
    T = len(pred_z_bins)
    bin_centers = z_lo + (np.arange(n_z) + 0.5) * (z_hi - z_lo) / n_z
    h_pred = bin_centers[np.clip(pred_z_bins, 0, n_z - 1)]
    h_gt = bin_centers[np.clip(gt_z_bins, 0, n_z - 1)]

    out_w, out_h = out_wh
    fig = plt.figure(figsize=(out_w / 100, out_h / 100), dpi=100)
    ax = fig.add_subplot(111)
    times = np.arange(T)
    ax.plot(times, h_gt, "-o", color="black", label="GT", markersize=4, linewidth=1.5)
    ax.plot(times, h_pred, "-o", color="#1f77b4", label="Pred", markersize=4, linewidth=1.5)
    ax.set_xlim(-0.5, T - 0.5)
    ax.set_ylim(z_lo, z_hi)
    ax.set_xlabel("timestep", fontsize=9)
    ax.set_ylabel("height (m)", fontsize=9)
    ax.tick_params(labelsize=8)
    ax.grid(True, alpha=0.25)
    ax.legend(loc="upper right", fontsize=8, frameon=False)
    fig.tight_layout(pad=0.4)
    fig.canvas.draw()
    rgb = np.asarray(fig.canvas.buffer_rgba())[:, :, :3]
    plt.close(fig)
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    if bgr.shape[1] != out_w or bgr.shape[0] != out_h:
        bgr = cv2.resize(bgr, (out_w, out_h), interpolation=cv2.INTER_AREA)
    return bgr