oscillode/blender_maxwell/node_trees/maxwell_sim_nodes/sockets/base.py

import typing as typ
import typing_extensions as typx
import functools

import bpy

import pydantic as pyd
import sympy as sp
import sympy.physics.units as spu
from .. import contracts as ct

class MaxwellSimSocket(bpy.types.NodeSocket):
	# Fundamentals
	socket_type: ct.SocketType
	bl_label: str
	
	# Style
	display_shape: typx.Literal[
		"CIRCLE", "SQUARE", "DIAMOND", "CIRCLE_DOT", "SQUARE_DOT",
		"DIAMOND_DOT",
	]
	## We use the following conventions for shapes:
	## - CIRCLE: Single Value.
	## - SQUARE: Container of Value.
	## - DIAMOND: Pointer Value.
	## - +DOT: Uses Units
	socket_color: tuple
	
	# Options
	#link_limit: int = 0
	use_units: bool = False
	use_prelock: bool = False
	
	# Computed
	bl_idname: str
	
	####################
	# - Initialization
	####################
	def __init_subclass__(cls, **kwargs: typ.Any):
		super().__init_subclass__(**kwargs)  ## Yucky superclass setup.
		
		# Setup Blender ID for Node
		if not hasattr(cls, "socket_type"):
			msg = f"Socket class {cls} does not define 'socket_type'"
			raise ValueError(msg)
		cls.bl_idname = str(cls.socket_type.value)
		
		# Setup Locked Property for Node
		cls.__annotations__["locked"] = bpy.props.BoolProperty(
			name="Locked State",
			description="The lock-state of a particular socket, which determines the socket's user editability",
			default=False,
		)
		
		# Setup Style
		cls.socket_color = ct.SOCKET_COLORS[cls.socket_type]
		cls.socket_shape = ct.SOCKET_SHAPES[cls.socket_type]
		
		# Setup List
		cls.__annotations__["is_list"] = bpy.props.BoolProperty(
			name="Is List",
			description="Whether or not a particular socket is a list type socket",
			default=False,
			update=lambda self, context: self.sync_is_list(context)
		)
		
		# Configure Use of Units
		if cls.use_units:
			# Set Shape :)
			cls.socket_shape += "_DOT"
			
			if not (socket_units := ct.SOCKET_UNITS.get(cls.socket_type)):
				msg = "Tried to `use_units` on {cls.bl_idname} socket, but `SocketType` has no units defined in `contracts.SOCKET_UNITS`"
				raise RuntimeError(msg)
			
			# Current Unit
			cls.__annotations__["active_unit"] = bpy.props.EnumProperty(
				name="Unit",
				description="Choose a unit",
				items=[
					(unit_name, str(unit_value), str(unit_value))
					for unit_name, unit_value in socket_units["values"].items()
				],
				default=socket_units["default"],
				update=lambda self, context: self.sync_unit_change(),
			)
			
			# Previous Unit (for conversion)
			cls.__annotations__["prev_active_unit"] = bpy.props.StringProperty(
				default=socket_units["default"],
			)
	
	####################
	# - Action Chain
	####################
	def trigger_action(
		self,
		action: typx.Literal["enable_lock", "disable_lock", "value_changed", "show_preview", "show_plot"],
	) -> None:
		"""Called whenever the socket's output value has changed.
		
		This also invalidates any of the socket's caches.
		
		When called on an input node, the containing node's
		`trigger_action` method will be called with this socket.
		
		When called on a linked output node, the linked socket's
		`trigger_action` method will be called.
		"""
		# Forwards Chains
		if action in {"value_changed"}:
			## Input Socket
			if not self.is_output:
				self.node.trigger_action(action, socket_name=self.name)
			
			## Linked Output Socket
			elif self.is_output and self.is_linked:
				for link in self.links:
					link.to_socket.trigger_action(action)
		
		# Backwards Chains
		elif action in {"enable_lock", "disable_lock", "show_preview", "show_plot"}:
			if action == "enable_lock":
				self.locked = True
			
			if action == "disable_lock":
				self.locked = False
			
			## Output Socket
			if self.is_output:
				self.node.trigger_action(action, socket_name=self.name)
			
			## Linked Input Socket
			elif not self.is_output and self.is_linked:
				for link in self.links:
					link.from_socket.trigger_action(action)
	
	####################
	# - Action Chain: Event Handlers
	####################
	def sync_is_list(self, context: bpy.types.Context):
		"""Called when the "is_list_ property has been updated.
		"""
		if self.is_list:
			if self.use_units:
				self.display_shape = "SQUARE_DOT"
			else:
				self.display_shape = "SQUARE"
		
		self.trigger_action("value_changed")
	
	def sync_prop(self, prop_name: str, context: bpy.types.Context):
		"""Called when a property has been updated.
		"""
		if not hasattr(self, prop_name):
			msg = f"Property {prop_name} not defined on socket {self}"
			raise RuntimeError(msg)
		
		self.trigger_action("value_changed")
		
	def sync_link_added(self, link) -> bool:
		"""Called when a link has been added to this (input) socket.
		
		Returns a bool, whether or not the socket consents to the link change.
		"""
		if self.locked: return False
		if self.is_output:
			msg = f"Tried to sync 'link add' on output socket"
			raise RuntimeError(msg)
		
		self.trigger_action("value_changed")
		
		return True
	
	def sync_link_removed(self, from_socket) -> bool:
		"""Called when a link has been removed from this (input) socket.
		
		Returns a bool, whether or not the socket consents to the link change.
		"""
		if self.locked: return False
		if self.is_output:
			msg = f"Tried to sync 'link add' on output socket"
			raise RuntimeError(msg)
		
		self.trigger_action("value_changed")
		
		return True
	
	####################
	# - Data Chain
	####################
	@property
	def value(self) -> typ.Any:
		raise NotImplementedError
	@value.setter
	def value(self, value: typ.Any) -> None:
		raise NotImplementedError
	
	@property
	def value_list(self) -> typ.Any:
		return [self.value]
	@value_list.setter
	def value_list(self, value: typ.Any) -> None:
		raise NotImplementedError
	
	def value_as_unit_system(
		self,
		unit_system: dict,
		dimensionless: bool = True
	) -> typ.Any:
		## TODO: Caching could speed this boi up quite a bit
		
		unit_system_unit = unit_system[self.socket_type]
		return spu.convert_to(
			self.value,
			unit_system_unit,
		) / unit_system_unit
	
	@property
	def lazy_value(self) -> None:
		raise NotImplementedError
	@lazy_value.setter
	def lazy_value(self, lazy_value: typ.Any) -> None:
		raise NotImplementedError
	
	@property
	def lazy_value_list(self) -> typ.Any:
		return [self.lazy_value]
	@lazy_value_list.setter
	def lazy_value_list(self, value: typ.Any) -> None:
		raise NotImplementedError
	
	@property
	def capabilities(self) -> None:
		raise NotImplementedError
	
	def _compute_data(
		self,
		kind: ct.DataFlowKind = ct.DataFlowKind.Value,
	) -> typ.Any:
		"""Computes the internal data of this socket, ONLY.
		
		**NOTE**: Low-level method. Use `compute_data` instead.
		"""
		if kind == ct.DataFlowKind.Value:
			if self.is_list: return self.value_list
			else: return self.value
		elif kind == ct.DataFlowKind.LazyValue:
			if self.is_list: return self.lazy_value_list
			else: return self.lazy_value
			return self.lazy_value
		elif kind == ct.DataFlowKind.Capabilities:
			return self.capabilities
		
		return None
	
	def compute_data(
		self,
		kind: ct.DataFlowKind = ct.DataFlowKind.Value,
	):
		"""Computes the value of this socket, including all relevant factors:
			- If input socket, and unlinked, compute internal data.
			- If input socket, and linked, compute linked socket data.
			- If output socket, ask node for data.
		"""
		# Compute Output Socket
		## List-like sockets guarantee that a list of a thing is passed.
		if self.is_output:
			res = self.node.compute_output(self.name, kind=kind)
			if self.is_list and not isinstance(res, list): return [res]
			return res
		
		# Compute Input Socket
		## Unlinked: Retrieve Socket Value
		if not self.is_linked: return self._compute_data(kind)
		
		## Linked: Compute Output of Linked Sockets
		linked_values = [
			link.from_socket.compute_data(kind)
			for link in self.links
		]
		
		## Return Single Value / List of Values
		if len(linked_values) == 1: return linked_values[0]
		return linked_values
	
	####################
	# - Unit Properties
	####################
	@functools.cached_property
	def possible_units(self) -> dict[str, sp.Expr]:
		if not self.use_units:
			msg = "Tried to get possible units for socket {self}, but socket doesn't `use_units`"
			raise ValueError(msg)
		
		return ct.SOCKET_UNITS[
			self.socket_type
		]["values"]
	
	@property
	def unit(self) -> sp.Expr:
		return self.possible_units[self.active_unit]
	
	@property
	def prev_unit(self) -> sp.Expr:
		return self.possible_units[self.prev_active_unit]
	
	@unit.setter
	def unit(self, value: str | sp.Expr) -> None:
		# Retrieve Unit by String
		if isinstance(value, str) and value in self.possible_units:
			self.active_unit = self.possible_units[value]
			return
		
		# Retrieve =1 Matching Unit Name
		matching_unit_names = [
			unit_name
			for unit_name, unit_sympy in self.possible_units.items()
			if value == unit_sympy
		]
		if len(matching_unit_names) == 0:
			msg = f"Tried to set unit for socket {self} with value {value}, but it is not one of possible units {''.join(possible.units.values())} for this socket (as defined in `contracts.SOCKET_UNITS`)"
			raise ValueError(msg)
		
		if len(matching_unit_names) > 1:
			msg = f"Tried to set unit for socket {self} with value {value}, but multiple possible matching units {''.join(possible.units.values())} for this socket (as defined in `contracts.SOCKET_UNITS`); there may only be one"
			raise RuntimeError(msg)
		
		self.active_unit = matching_unit_names[0]
	
	def sync_unit_change(self) -> None:
		"""In unit-aware sockets, the internal `value()` property multiplies the Blender property value by the current active unit.
		
		When the unit is changed, `value()` will display the old scalar with the new unit.
		To fix this, we need to update the scalar to use the new unit.
		
		Can be overridden if more specific logic is required.
		"""
		
		prev_value = self.value / self.unit * self.prev_unit
		## After changing units, self.value is expressed in the wrong unit.
		## - Therefore, we removing the new unit, and re-add the prev unit.
		## - Using only self.value avoids implementation-specific details.
		
		self.value = spu.convert_to(
			prev_value,
			self.unit
		)  ## Now, the unit conversion can be done correctly.
		
		self.prev_active_unit = self.active_unit
	
	####################
	# - Style
	####################
	def draw_color(
		self,
		context: bpy.types.Context,
		node: bpy.types.Node,
	) -> ct.BLColorRGBA:
		"""Color of the socket icon, when embedded in a node.
		"""
		return self.socket_color
	
	@classmethod
	def draw_color_simple(cls) -> ct.BLColorRGBA:
		"""Fallback color of the socket icon (ex.when not embedded in a node).
		"""
		return cls.socket_color
	
	####################
	# - UI Methods
	####################
	def draw(
		self,
		context: bpy.types.Context,
		layout: bpy.types.UILayout,
		node: bpy.types.Node,
		text: str,
	) -> None:
		"""Called by Blender to draw the socket UI.
		"""
		
		if self.is_output:
			self.draw_output(context, layout, node, text)
		else:
			self.draw_input(context, layout, node, text)
	
	def draw_prelock(
		self,
		context: bpy.types.Context,
		col: bpy.types.UILayout,
		node: bpy.types.Node,
		text: str,
	) -> None:
		pass
	
	def draw_input(
		self,
		context: bpy.types.Context,
		layout: bpy.types.UILayout,
		node: bpy.types.Node,
		text: str,
	) -> None:
		"""Draws the socket UI, when the socket is an input socket.
		"""
		col = layout.column(align=False)
		
		# Label Row
		row = col.row(align=False)
		if self.locked: row.enabled = False
		
		## Linked Label
		if self.is_linked:
			row.label(text=text)
			return
		
		## User Label Row (incl. Units)
		if self.use_units:
			split = row.split(factor=0.6, align=True)
			
			_row = split.row(align=True)
			self.draw_label_row(_row, text)
		
			_col = split.column(align=True)
			_col.prop(self, "active_unit", text="")
		else:
			self.draw_label_row(row, text)
		
		# Prelock Row
		row = col.row(align=False)
		if self.use_prelock:
			_col = row.column(align=False)
			_col.enabled = True
			self.draw_prelock(context, _col, node, text)
			
			if self.locked:
				row = col.row(align=False)
				row.enabled = False
		else:
			if self.locked: row.enabled = False
		
		# Value Column(s)
		col = row.column(align=True)
		if self.is_list:
			self.draw_value_list(col)
		else:
			self.draw_value(col)
	
	def draw_output(
		self,
		context: bpy.types.Context,
		layout: bpy.types.UILayout,
		node: bpy.types.Node,
		text: str,
	) -> None:
		"""Draws the socket UI, when the socket is an output socket.
		"""
		layout.label(text=text)
	
	####################
	# - UI Methods
	####################
	def draw_label_row(
		self,
		row: bpy.types.UILayout,
		text: str,
	) -> None:
		"""Called to draw the label row (same height as socket shape).
		
		Can be overridden.
		"""
		row.label(text=text)
	
	def draw_value(self, col: bpy.types.UILayout) -> None:
		"""Called to draw the value column in unlinked input sockets.
		
		Can be overridden.
		"""
		pass
	
	def draw_value_list(self, col: bpy.types.UILayout) -> None:
		"""Called to draw the value list column in unlinked input sockets.
		
		Can be overridden.
		"""
		pass