StarPilot/system/ui/lib/shader_polygon.py

import os
import pyray as rl
import numpy as np
from raylib import rl as raw_rl
from dataclasses import dataclass
from typing import Any, Optional, cast
from openpilot.system.ui.lib.application import gui_app, GL_VERSION

USE_VERTEX_GRADIENTS = os.getenv("USE_VERTEX_GRADIENTS", "1") == "1"

MAX_GRADIENT_COLORS = 20  # includes stops as well


@dataclass
class Gradient:
  start: tuple[float, float]
  end: tuple[float, float]
  colors: list[rl.Color]
  stops: list[float]

  def __post_init__(self):
    if len(self.colors) > MAX_GRADIENT_COLORS:
      self.colors = self.colors[:MAX_GRADIENT_COLORS]
      print(f"Warning: Gradient colors truncated to {MAX_GRADIENT_COLORS} entries")

    if len(self.stops) > MAX_GRADIENT_COLORS:
      self.stops = self.stops[:MAX_GRADIENT_COLORS]
      print(f"Warning: Gradient stops truncated to {MAX_GRADIENT_COLORS} entries")

    if not len(self.stops):
      color_count = min(len(self.colors), MAX_GRADIENT_COLORS)
      self.stops = [i / max(1, color_count - 1) for i in range(color_count)]


FRAGMENT_SHADER = GL_VERSION + """
in vec2 fragTexCoord;
out vec4 finalColor;

uniform vec4 fillColor;

// Gradient line defined in *screen pixels*
uniform int useGradient;
uniform vec2 gradientStart;  // e.g. vec2(0, 0)
uniform vec2 gradientEnd;    // e.g. vec2(0, screenHeight)
uniform vec4 gradientColors[20];
uniform float gradientStops[20];
uniform int gradientColorCount;

vec4 getGradientColor(vec2 p) {
  // Compute t from screen-space position
  vec2 d = gradientStart - gradientEnd;
  float len2 = max(dot(d, d), 1e-6);
  float t = clamp(dot(p - gradientEnd, d) / len2, 0.0, 1.0);

  // Clamp to range
  float t0 = gradientStops[0];
  float tn = gradientStops[gradientColorCount-1];
  if (t <= t0) return gradientColors[0];
  if (t >= tn) return gradientColors[gradientColorCount-1];

  for (int i = 0; i < gradientColorCount - 1; i++) {
    float a = gradientStops[i];
    float b = gradientStops[i+1];
    if (t >= a && t <= b) {
      float k = (t - a) / max(b - a, 1e-6);
      return mix(gradientColors[i], gradientColors[i+1], k);
    }
  }

  return gradientColors[gradientColorCount-1];
}

void main() {
  // TODO: do proper antialiasing
  finalColor = useGradient == 1 ? getGradientColor(gl_FragCoord.xy) : fillColor;
}
"""

# Default vertex shader
VERTEX_SHADER = GL_VERSION + """
in vec3 vertexPosition;
in vec2 vertexTexCoord;
out vec2 fragTexCoord;
uniform mat4 mvp;

void main() {
  fragTexCoord = vertexTexCoord;
  gl_Position = mvp * vec4(vertexPosition, 1.0);
}
"""

UNIFORM_INT = rl.ShaderUniformDataType.SHADER_UNIFORM_INT
UNIFORM_FLOAT = rl.ShaderUniformDataType.SHADER_UNIFORM_FLOAT
UNIFORM_VEC2 = rl.ShaderUniformDataType.SHADER_UNIFORM_VEC2
UNIFORM_VEC4 = rl.ShaderUniformDataType.SHADER_UNIFORM_VEC4


class ShaderState:
  _instance: Any = None

  @classmethod
  def get_instance(cls):
    if cls._instance is None:
      cls._instance = cls()
    return cls._instance

  def __init__(self):
    if ShaderState._instance is not None:
      raise Exception("This class is a singleton. Use get_instance() instead.")

    self.initialized = False
    self.shader = None

    # Shader uniform locations
    self.locations = {
      'fillColor': None,
      'useGradient': None,
      'gradientStart': None,
      'gradientEnd': None,
      'gradientColors': None,
      'gradientStops': None,
      'gradientColorCount': None,
      'mvp': None,
    }

    # Pre-allocated FFI objects
    self.fill_color_ptr = rl.ffi.new("float[]", [0.0, 0.0, 0.0, 0.0])
    self.use_gradient_ptr = rl.ffi.new("int[]", [0])
    self.color_count_ptr = rl.ffi.new("int[]", [0])
    self.gradient_colors_ptr = rl.ffi.new("float[]", MAX_GRADIENT_COLORS * 4)
    self.gradient_stops_ptr = rl.ffi.new("float[]", MAX_GRADIENT_COLORS)

  def initialize(self):
    if self.initialized:
      return

    self.shader = rl.load_shader_from_memory(VERTEX_SHADER, FRAGMENT_SHADER)

    # Cache all uniform locations
    for uniform in self.locations.keys():
      self.locations[uniform] = rl.get_shader_location(self.shader, uniform)

    # Orthographic MVP (origin top-left)
    proj = rl.matrix_ortho(0, gui_app.width, gui_app.height, 0, -1, 1)
    rl.set_shader_value_matrix(self.shader, self.locations['mvp'], proj)

    self.initialized = True

  def cleanup(self):
    if not self.initialized:
      return
    if self.shader:
      rl.unload_shader(self.shader)
      self.shader = None

    self.initialized = False


def _configure_shader_color(state: ShaderState, color: Optional[rl.Color],
                            gradient: Gradient | None, origin_rect: rl.Rectangle):
  assert (color is not None) != (gradient is not None), "Either color or gradient must be provided"

  use_gradient = 1 if (gradient is not None and len(gradient.colors) >= 1) else 0
  state.use_gradient_ptr[0] = use_gradient
  rl.set_shader_value(state.shader, state.locations['useGradient'], state.use_gradient_ptr, UNIFORM_INT)

  if use_gradient:
    gradient = cast(Gradient, gradient)
    state.color_count_ptr[0] = len(gradient.colors)
    for i in range(len(gradient.colors)):
      c = gradient.colors[i]
      base = i * 4
      state.gradient_colors_ptr[base:base + 4] = [c.r / 255.0, c.g / 255.0, c.b / 255.0, c.a / 255.0]
    rl.set_shader_value_v(state.shader, state.locations['gradientColors'], state.gradient_colors_ptr, UNIFORM_VEC4, len(gradient.colors))

    for i in range(len(gradient.stops)):
      s = float(gradient.stops[i])
      state.gradient_stops_ptr[i] = 0.0 if s < 0.0 else 1.0 if s > 1.0 else s
    rl.set_shader_value_v(state.shader, state.locations['gradientStops'], state.gradient_stops_ptr, UNIFORM_FLOAT, len(gradient.stops))
    rl.set_shader_value(state.shader, state.locations['gradientColorCount'], state.color_count_ptr, UNIFORM_INT)

    # Map normalized start/end to screen pixels
    start_vec = rl.Vector2(origin_rect.x + gradient.start[0] * origin_rect.width, origin_rect.y + gradient.start[1] * origin_rect.height)
    end_vec = rl.Vector2(origin_rect.x + gradient.end[0] * origin_rect.width, origin_rect.y + gradient.end[1] * origin_rect.height)
    rl.set_shader_value(state.shader, state.locations['gradientStart'], start_vec, UNIFORM_VEC2)
    rl.set_shader_value(state.shader, state.locations['gradientEnd'], end_vec, UNIFORM_VEC2)
  else:
    color = color or rl.WHITE
    state.fill_color_ptr[0:4] = [color.r / 255.0, color.g / 255.0, color.b / 255.0, color.a / 255.0]
    rl.set_shader_value(state.shader, state.locations['fillColor'], state.fill_color_ptr, UNIFORM_VEC4)


def triangulate(pts: np.ndarray) -> np.ndarray:
  """Only supports simple polygons with two chains (ribbon)."""

  # interleave points to produce a triangle strip
  if len(pts) % 2 != 0:
    pts = pts[:-1]

  half = len(pts) // 2
  tri_strip = np.empty(pts.shape, dtype=np.float32)
  tri_strip[0::2] = pts[:half]
  tri_strip[1::2] = pts[half:][::-1]
  return tri_strip


def _draw_strip_vertex_gradient(tri_strip: np.ndarray, gradient: Gradient, origin_rect: rl.Rectangle) -> None:
  n = len(tri_strip)
  if n < 3 or not gradient.colors:
    return

  start_x = origin_rect.x + gradient.start[0] * origin_rect.width
  start_y = origin_rect.y + gradient.start[1] * origin_rect.height
  end_x = origin_rect.x + gradient.end[0] * origin_rect.width
  end_y = origin_rect.y + gradient.end[1] * origin_rect.height

  dx = start_x - end_x
  dy = start_y - end_y
  len2 = max(dx * dx + dy * dy, 1e-6)

  vx = tri_strip[:, 0] - end_x
  vy = tri_strip[:, 1] - end_y
  t = np.clip((vx * dx + vy * dy) / len2, 0.0, 1.0)

  colors = gradient.colors
  stops = gradient.stops if len(gradient.stops) == len(colors) else [i / max(1, len(colors) - 1) for i in range(len(colors))]
  r_stops = [c.r for c in colors]
  g_stops = [c.g for c in colors]
  b_stops = [c.b for c in colors]
  a_stops = [c.a for c in colors]

  r = np.interp(t, stops, r_stops).astype(np.uint8)
  g = np.interp(t, stops, g_stops).astype(np.uint8)
  b = np.interp(t, stops, b_stops).astype(np.uint8)
  a = np.interp(t, stops, a_stops).astype(np.uint8)

  rl_begin = raw_rl.rlBegin
  rl_end = raw_rl.rlEnd
  rl_color = raw_rl.rlColor4ub
  rl_vertex = raw_rl.rlVertex2f

  rl_begin(4)  # RL_TRIANGLES
  for i in range(2, n):
    if i % 2 == 0:
      i0, i1, i2 = i, i - 2, i - 1
    else:
      i0, i1, i2 = i, i - 1, i - 2
    rl_color(r[i0], g[i0], b[i0], a[i0])
    rl_vertex(tri_strip[i0, 0], tri_strip[i0, 1])
    rl_color(r[i1], g[i1], b[i1], a[i1])
    rl_vertex(tri_strip[i1, 0], tri_strip[i1, 1])
    rl_color(r[i2], g[i2], b[i2], a[i2])
    rl_vertex(tri_strip[i2, 0], tri_strip[i2, 1])
  rl_end()


def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
                 color: Optional[rl.Color] = None, gradient: Gradient | None = None):

  """
  Draw a ribbon polygon (two chains) with a triangle strip and gradient.
  - Input must be [L0..Lk-1, Rk-1..R0], even count, no crossings/holes.
  """
  if len(points) < 3:
    return

  pts = np.asarray(points)
  assert pts.ndim == 2 and pts.shape[1] == 2, "points must be (N,2)"

  # Triangulate via interleaving
  tri_strip = triangulate(pts)
  if len(tri_strip) < 3:
    return

  if gradient is None:
    vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)
    rl.draw_triangle_strip(vertices, len(tri_strip), color or rl.WHITE)
    return

  if USE_VERTEX_GRADIENTS:
    _draw_strip_vertex_gradient(tri_strip, gradient, origin_rect)
    return

  vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)
  state = ShaderState.get_instance()
  state.initialize()

  _configure_shader_color(state, color, gradient, origin_rect)

  # Draw strip, color here doesn't matter
  rl.begin_shader_mode(state.shader)
  rl.draw_triangle_strip(vertices, len(tri_strip), rl.WHITE)
  rl.end_shader_mode()


def cleanup_shader_resources():
  state = ShaderState.get_instance()
  state.cleanup()