oscillode/blender_maxwell/node_trees/maxwell_sim_nodes/contracts.py

import typing as typ
import typing_extensions as pytypes_ext
import enum

import sympy as sp

sp.printing.str.StrPrinter._default_settings['abbrev'] = True
## When we str() a unit expression, use abbrevied units.

import sympy.physics.units as spu
import pydantic as pyd
import bpy

from ...utils.blender_type_enum import (
	BlenderTypeEnum, append_cls_name_to_values, wrap_values_in_MT
)

####################
# - String Types
####################
BlenderColorRGB = tuple[float, float, float, float]
BlenderID = pytypes_ext.Annotated[str, pyd.StringConstraints(
	pattern=r'^[A-Z_]+$',
)]

# Socket ID
SocketName = pytypes_ext.Annotated[str, pyd.StringConstraints(
	pattern=r'^[a-zA-Z0-9_]+$',
)]
BLSocketName = pytypes_ext.Annotated[str, pyd.StringConstraints(
	pattern=r'^[a-zA-Z0-9_]+$',
)]

# Socket ID
PresetID = pytypes_ext.Annotated[str, pyd.StringConstraints(
	pattern=r'^[A-Z_]+$',
)]

####################
# - Sympy Expression Typing
####################
ALL_UNIT_SYMBOLS = {
	unit
	for unit in spu.__dict__.values()
	if isinstance(unit, spu.Quantity)
}
def has_units(expr: sp.Expr):
	return any(
		symbol in ALL_UNIT_SYMBOLS
		for symbol in expr.atoms(sp.Symbol)
	)
def is_exactly_expressed_as_unit(expr: sp.Expr, unit) -> bool:
	#try:
	converted_expr = expr / unit
	
	return (
		converted_expr.is_number
		and not converted_expr.has(spu.Quantity)
	)

####################
# - Icon Types
####################
class Icon(BlenderTypeEnum):
	MaxwellSimTree = "MOD_SIMPLEDEFORM"

####################
# - Tree Types
####################
@append_cls_name_to_values
class TreeType(BlenderTypeEnum):
	MaxwellSim = enum.auto()

####################
# - Socket Types
####################
@append_cls_name_to_values
class SocketType(BlenderTypeEnum):
	# Base
	Any = enum.auto()
	Text = enum.auto()
	FilePath = enum.auto()
	
	# Number
	IntegerNumber = enum.auto()
	RationalNumber = enum.auto()
	RealNumber = enum.auto()
	ComplexNumber = enum.auto()
	
	# Vector
	Real2DVector = enum.auto()
	Complex2DVector = enum.auto()
	
	Real3DVector = enum.auto()
	Complex3DVector = enum.auto()
	
	# Physical
	PhysicalTime = enum.auto()
	
	PhysicalAngle = enum.auto()
	
	PhysicalLength = enum.auto()
	PhysicalArea = enum.auto()
	PhysicalVolume = enum.auto()
	
	PhysicalMass = enum.auto()
	
	PhysicalSpeed = enum.auto()
	PhysicalAccel = enum.auto()
	PhysicalForce = enum.auto()
	
	PhysicalPol = enum.auto()
	
	PhysicalFreq = enum.auto()
	PhysicalSpecPowerDist = enum.auto()
	PhysicalSpecRelPermDist = enum.auto()
	
	# Blender
	BlenderObject = enum.auto()
	BlenderCollection = enum.auto()
	
	BlenderImage = enum.auto()
	BlenderVolume = enum.auto()
	
	BlenderGeoNodes = enum.auto()
	BlenderText = enum.auto()
	
	# Maxwell
	MaxwellSource = enum.auto()
	MaxwellTemporalShape = enum.auto()
	
	MaxwellMedium = enum.auto()
	MaxwellMediumNonLinearity = enum.auto()
	
	MaxwellStructure = enum.auto()
	
	MaxwellBoundBox = enum.auto()
	MaxwellBoundFace = enum.auto()
	
	MaxwellMonitor = enum.auto()

	MaxwellFDTDSim = enum.auto()
	MaxwellSimGrid = enum.auto()
	MaxwellSimGridAxis = enum.auto()

SocketType_to_units = {
	SocketType.PhysicalTime: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	
	SocketType.PhysicalAngle: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	
	SocketType.PhysicalLength: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	SocketType.PhysicalArea: {
		"default": "UM_SQ",
		"values": {
			"PM_SQ": spu.pm**2,
			"A_SQ": spu.angstrom**2,
			"NM_SQ": spu.nm**2,
			"UM_SQ": spu.um**2,
			"MM_SQ": spu.mm**2,
			"CM_SQ": spu.cm**2,
			"M_SQ": spu.m**2,
		},
	},
	SocketType.PhysicalVolume: {
		"default": "UM_CB",
		"values": {
			"PM_CB": spu.pm**3,
			"A_CB": spu.angstrom**3,
			"NM_CB": spu.nm**3,
			"UM_CB": spu.um**3,
			"MM_CB": spu.mm**3,
			"CM_CB": spu.cm**3,
			"M_CB": spu.m**3,
			"ML": spu.milliliter,
			"L": spu.liter,
		},
	},
	
	SocketType.PhysicalMass: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	
	SocketType.PhysicalSpeed: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	SocketType.PhysicalAccel: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	SocketType.PhysicalForce: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	
	SocketType.PhysicalPol: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	
	SocketType.PhysicalFreq: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	SocketType.PhysicalSpecPowerDist: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
	SocketType.PhysicalSpecRelPermDist: {
		"default": "UM",
		"values": {
			"UM": spu.um,
		},
	},
}

SocketType_to_color = {
	SocketType.Any: (0.5, 0.5, 0.5, 1.0),
	SocketType.Text: (0.5, 0.5, 0.5, 1.0),
	SocketType.FilePath: (0.5, 0.5, 0.5, 1.0),
	
	# Mathematical
	SocketType.IntegerNumber: (0.5, 0.5, 0.5, 1.0),
	SocketType.RationalNumber: (0.5, 0.5, 0.5, 1.0),
	SocketType.RealNumber: (0.5, 0.5, 0.5, 1.0),
	SocketType.ComplexNumber: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.Real2DVector: (0.5, 0.5, 0.5, 1.0),
	SocketType.Complex2DVector: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.Real3DVector: (0.5, 0.5, 0.5, 1.0),
	SocketType.Complex3DVector: (0.5, 0.5, 0.5, 1.0),
	
	# Physical
	SocketType.PhysicalTime: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalAngle: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalLength: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalArea: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalVolume: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalMass: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalSpeed: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalAccel: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalForce: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalPol: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.PhysicalFreq: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalSpecPowerDist: (0.5, 0.5, 0.5, 1.0),
	SocketType.PhysicalSpecRelPermDist: (0.5, 0.5, 0.5, 1.0),
	
	# Blender
	SocketType.BlenderObject: (0.5, 0.5, 0.5, 1.0),
	SocketType.BlenderCollection: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.BlenderImage: (0.5, 0.5, 0.5, 1.0),
	SocketType.BlenderVolume: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.BlenderGeoNodes: (0.5, 0.5, 0.5, 1.0),
	SocketType.BlenderText: (0.5, 0.5, 0.5, 1.0),
	
	# Maxwell
	SocketType.MaxwellSource: (0.5, 0.5, 0.5, 1.0),
	SocketType.MaxwellTemporalShape: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.MaxwellMedium: (0.5, 0.5, 0.5, 1.0),
	SocketType.MaxwellMediumNonLinearity: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.MaxwellStructure: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.MaxwellBoundBox: (0.5, 0.5, 0.5, 1.0),
	SocketType.MaxwellBoundFace: (0.5, 0.5, 0.5, 1.0),
	
	SocketType.MaxwellMonitor: (0.5, 0.5, 0.5, 1.0),

	SocketType.MaxwellFDTDSim: (0.5, 0.5, 0.5, 1.0),
	SocketType.MaxwellSimGrid: (0.5, 0.5, 0.5, 1.0),
	SocketType.MaxwellSimGridAxis: (0.5, 0.5, 0.5, 1.0),
}

####################
# - Node Types
####################
@append_cls_name_to_values
class NodeType(BlenderTypeEnum):
	KitchenSink = enum.auto()
	
	# Inputs
	## Inputs / Scene
	Time = enum.auto()
	UnitSystem = enum.auto()
	
	## Inputs / Parameters
	NumberParameter = enum.auto()
	PhysicalParameter = enum.auto()
	
	## Inputs / Constants
	ScientificConstant = enum.auto()
	NumberConstant = enum.auto()
	PhysicalConstant = enum.auto()
	BlenderConstant = enum.auto()
	
	## Inputs / Lists
	RealList = enum.auto()
	ComplexList = enum.auto()
	
	## Inputs / 
	InputFile = enum.auto()
	
	
	# Outputs
	## Outputs / Viewers
	ValueViewer = enum.auto()
	ConsoleViewer = enum.auto()
	
	## Outputs / Exporters
	JSONFileExporter = enum.auto()
	
	
	# Sources
	## Sources / Temporal Shapes
	GaussianPulseTemporalShape = enum.auto()
	ContinuousWaveTemporalShape = enum.auto()
	ListTemporalShape = enum.auto()
	
	## Sources /
	PointDipoleSource = enum.auto()
	UniformCurrentSource = enum.auto()
	PlaneWaveSource = enum.auto()
	ModeSource = enum.auto()
	GaussianBeamSource = enum.auto()
	AstigmaticGaussianBeamSource = enum.auto()
	TFSFSource = enum.auto()
	
	EHEquivalenceSource = enum.auto()
	EHSource = enum.auto()
	
	# Mediums
	LibraryMedium = enum.auto()
	
	PECMedium = enum.auto()
	IsotropicMedium = enum.auto()
	AnisotropicMedium = enum.auto()
	
	TripleSellmeierMedium = enum.auto()
	SellmeierMedium = enum.auto()
	PoleResidueMedium = enum.auto()
	DrudeMedium = enum.auto()
	DrudeLorentzMedium = enum.auto()
	DebyeMedium = enum.auto()
	
	## Mediums / Non-Linearities
	AddNonLinearity = enum.auto()
	ChiThreeSusceptibilityNonLinearity = enum.auto()
	TwoPhotonAbsorptionNonLinearity = enum.auto()
	KerrNonLinearity = enum.auto()
	
	# Structures
	ObjectStructure = enum.auto()
	GeoNodesStructure = enum.auto()
	ScriptedStructure = enum.auto()
	
	## Structures / Primitives
	BoxStructure = enum.auto()
	SphereStructure = enum.auto()
	CylinderStructure = enum.auto()
	
	
	# Bounds
	BoundBox = enum.auto()
	
	## Bounds / Bound Faces
	PMLBoundFace = enum.auto()
	PECBoundFace = enum.auto()
	PMCBoundFace = enum.auto()
	
	BlochBoundFace = enum.auto()
	PeriodicBoundFace = enum.auto()
	AbsorbingBoundFace = enum.auto()
	
	
	# Monitors
	EHFieldMonitor = enum.auto()
	FieldPowerFluxMonitor = enum.auto()
	EpsilonTensorMonitor = enum.auto()
	DiffractionMonitor = enum.auto()
	
	## Monitors / Near-Field Projections
	CartesianNearFieldProjectionMonitor = enum.auto()
	ObservationAngleNearFieldProjectionMonitor = enum.auto()
	KSpaceNearFieldProjectionMonitor = enum.auto()
	
	
	# Sims
	SimGrid = enum.auto()
	
	## Sims / Sim Grid Axis
	AutomaticSimGridAxis = enum.auto()
	ManualSimGridAxis = enum.auto()
	UniformSimGridAxis = enum.auto()
	ArraySimGridAxis = enum.auto()
	
	## Sim /
	FDTDSim = enum.auto()
	
	
	# Utilities
	Math = enum.auto()
	
	## Utilities / Operations
	ArrayOperation = enum.auto()

####################
# - Node Category Types
####################
@wrap_values_in_MT
class NodeCategory(BlenderTypeEnum):
	MAXWELLSIM = enum.auto()
	
	# Inputs/
	MAXWELLSIM_INPUTS = enum.auto()
	MAXWELLSIM_INPUTS_SCENE = enum.auto()
	MAXWELLSIM_INPUTS_PARAMETERS = enum.auto()
	MAXWELLSIM_INPUTS_CONSTANTS = enum.auto()
	MAXWELLSIM_INPUTS_LISTS = enum.auto()
	
	# Outputs/
	MAXWELLSIM_OUTPUTS = enum.auto()
	MAXWELLSIM_OUTPUTS_VIEWERS = enum.auto()
	MAXWELLSIM_OUTPUTS_EXPORTERS = enum.auto()
	MAXWELLSIM_OUTPUTS_PLOTTERS = enum.auto()
	
	# Sources/
	MAXWELLSIM_SOURCES = enum.auto()
	MAXWELLSIM_SOURCES_TEMPORALSHAPES = enum.auto()
	
	# Mediums/
	MAXWELLSIM_MEDIUMS = enum.auto()
	MAXWELLSIM_MEDIUMS_NONLINEARITIES = enum.auto()
	
	# Structures/
	MAXWELLSIM_STRUCTURES = enum.auto()
	MAXWELLSIM_STRUCTURES_PRIMITIVES = enum.auto()
	
	# Bounds/
	MAXWELLSIM_BOUNDS = enum.auto()
	MAXWELLSIM_BOUNDS_BOUNDFACES = enum.auto()
	
	# Monitors/
	MAXWELLSIM_MONITORS = enum.auto()
	MAXWELLSIM_MONITORS_NEARFIELDPROJECTIONS = enum.auto()
	
	# Simulations/
	MAXWELLSIM_SIMS = enum.auto()
	MAXWELLSIM_SIMGRIDAXES = enum.auto()
	
	# Utilities/
	MAXWELLSIM_UTILITIES = enum.auto()
	MAXWELLSIM_UTILITIES_OPERATIONS = enum.auto()
	
	@classmethod
	def get_tree(cls):
		## TODO: Refactor
		syllable_categories = [
			node_category.value.split("_")
			for node_category in cls
			if node_category.value != "MAXWELLSIM"
		]
		
		category_tree = {}
		for syllable_category in syllable_categories:
			# Set Current Subtree to Root
			current_category_subtree = category_tree
			
			for i, syllable in enumerate(syllable_category):
				# Create New Category Subtree and/or Step to Subtree
				if syllable not in current_category_subtree:
					current_category_subtree[syllable] = {}
				current_category_subtree = current_category_subtree[syllable]
		
		return category_tree

NodeCategory_to_category_label = {
	# Inputs/
	NodeCategory.MAXWELLSIM_INPUTS: "Inputs",
	NodeCategory.MAXWELLSIM_INPUTS_SCENE: "Scene",
	NodeCategory.MAXWELLSIM_INPUTS_PARAMETERS: "Parameters",
	NodeCategory.MAXWELLSIM_INPUTS_CONSTANTS: "Constants",
	NodeCategory.MAXWELLSIM_INPUTS_LISTS: "Lists",
	
	# Outputs/
	NodeCategory.MAXWELLSIM_OUTPUTS: "Outputs",
	NodeCategory.MAXWELLSIM_OUTPUTS_VIEWERS: "Viewers",
	NodeCategory.MAXWELLSIM_OUTPUTS_EXPORTERS: "Exporters",
	NodeCategory.MAXWELLSIM_OUTPUTS_PLOTTERS: "Plotters",
	
	# Sources/
	NodeCategory.MAXWELLSIM_SOURCES: "Sources",
	NodeCategory.MAXWELLSIM_SOURCES_TEMPORALSHAPES: "Temporal Shapes",
	
	# Mediums/
	NodeCategory.MAXWELLSIM_MEDIUMS: "Mediums",
	NodeCategory.MAXWELLSIM_MEDIUMS_NONLINEARITIES: "Non-Linearities",
	
	# Structures/
	NodeCategory.MAXWELLSIM_STRUCTURES: "Structures",
	NodeCategory.MAXWELLSIM_STRUCTURES_PRIMITIVES: "Primitives",
	
	# Bounds/
	NodeCategory.MAXWELLSIM_BOUNDS: "Bounds",
	NodeCategory.MAXWELLSIM_BOUNDS_BOUNDFACES: "Bound Faces",
	
	# Monitors/
	NodeCategory.MAXWELLSIM_MONITORS: "Monitors",
	NodeCategory.MAXWELLSIM_MONITORS_NEARFIELDPROJECTIONS: "Near-Field Projections",
	
	# Simulations/
	NodeCategory.MAXWELLSIM_SIMS: "Simulations",
	NodeCategory.MAXWELLSIM_SIMGRIDAXES: "Sim Grid Axes",
	
	# Utilities/
	NodeCategory.MAXWELLSIM_UTILITIES: "Utilities",
	NodeCategory.MAXWELLSIM_UTILITIES_OPERATIONS: "Operations",
}



####################
# - Protocols
####################
class SocketDefProtocol(typ.Protocol):
	socket_type: SocketType
	label: str
	
	def init(self, bl_socket: bpy.types.NodeSocket) -> None:
		...

class PresetDef(pyd.BaseModel):
	label: str
	description: str
	values: dict[SocketName, typ.Any]

SocketReturnType = typ.TypeVar('SocketReturnType', covariant=True)
## - Covariance: If B subtypes A, then Container[B] subtypes Container[A].
## - This is absolutely what we want here.

#@typ.runtime_checkable
#class BLSocketProtocol(typ.Protocol):
#	socket_type: SocketType
#	socket_color: BlenderColorRGB
#	
#	bl_label: str
#	
#	compatible_types: dict[typ.Type, set[typ.Callable[[typ.Any], bool]]]
#	
#	def draw(
#		self,
#		context: bpy.types.Context,
#		layout: bpy.types.UILayout,
#		node: bpy.types.Node,
#		text: str,
#	) -> None:
#		...
#	
#	@property
#	def default_value(self) -> typ.Any:
#		...
#	@default_value.setter
#	def default_value(self, value: typ.Any) -> typ.Any:
#		...
#	

@typ.runtime_checkable
class NodeTypeProtocol(typ.Protocol):
	node_type: NodeType
	
	bl_label: str
	
	input_sockets: dict[SocketName, SocketDefProtocol]
	output_sockets: dict[SocketName, SocketDefProtocol]
	presets: dict[PresetID, PresetDef] | None
	
	# Built-In Blender Methods
	def init(self, context: bpy.types.Context) -> None:
		...
	
	def draw_buttons(
		self,
		context: bpy.types.Context,
		layout: bpy.types.UILayout,
	) -> None:
		...
	
	@classmethod
	def poll(cls, ntree: bpy.types.NodeTree) -> None:
		...
	
	# Socket Getters
	def g_input_bl_socket(
		self,
		input_socket_name: SocketName,
	) -> bpy.types.NodeSocket:
		...
	
	def g_output_bl_socket(
		self,
		output_socket_name: SocketName,
	) -> bpy.types.NodeSocket:
		...
	
	# Socket Methods
	def s_input_value(
		self,
		input_socket_name: SocketName,
		value: typ.Any
	) -> typ.Any:
		...
	
	# Data-Flow Methods
	def compute_input(
		self,
		input_socket_name: SocketName,
	) -> typ.Any:
		...
	def compute_output(
		self,
		output_socket_name: SocketName,
	) -> typ.Any:
		...