Atomizer/atomizer-dashboard/frontend/src/lib/canvas/schema.ts

/**
 * Canvas Schema - Type definitions for optimization workflow nodes
 */

export type NodeType =
  | 'model'
  | 'solver'
  | 'designVar'
  | 'extractor'
  | 'objective'
  | 'constraint'
  | 'algorithm'
  | 'surrogate';

export interface BaseNodeData {
  label: string;
  configured: boolean;
  errors?: string[];
}

export interface ModelNodeData extends BaseNodeData {
  type: 'model';
  filePath?: string;
  fileType?: 'prt' | 'fem' | 'sim';
}

export interface SolverNodeData extends BaseNodeData {
  type: 'solver';
  solverType?: 'SOL101' | 'SOL103' | 'SOL105' | 'SOL106' | 'SOL111' | 'SOL112';
}

export interface DesignVarNodeData extends BaseNodeData {
  type: 'designVar';
  expressionName?: string;
  minValue?: number;
  maxValue?: number;
  baseline?: number;
  unit?: string;
  enabled?: boolean;
  notes?: string;
}

// Extractor-specific config types
export interface ZernikeConfig {
  innerRadius?: number;
  outerRadius?: number;
  nModes?: number;
  filterLowOrders?: number;
  subcases?: string[];
  subcaseLabels?: Record<string, string>;
  referenceSubcase?: string;
  extractMethod?: 'extract_relative' | 'extract_rms' | 'extract_absolute';
}

export interface DisplacementConfig {
  subcase?: number;
  nodeSet?: string;
  component?: 'magnitude' | 'x' | 'y' | 'z';
}

export interface StressConfig {
  subcase?: number;
  elementSet?: string;
  stressType?: 'vonMises' | 'principal' | 'max_shear';
}

export interface MassConfig {
  source?: 'bdf' | 'expression';
  expressionName?: string;
}

export interface FrequencyConfig {
  modeNumber?: number;
}

export type ExtractorConfig = Record<string, unknown>;

export interface ExtractorNodeData extends BaseNodeData {
  type: 'extractor';
  extractorId?: string;
  extractorName?: string;
  extractorType?: 'zernike_opd' | 'displacement' | 'stress' | 'mass' | 'frequency';
  extractMethod?: string;
  config?: ExtractorConfig;
  // Zernike-specific (for quick access)
  innerRadius?: number;
  nModes?: number;
  subcases?: string[];
  // Output mapping
  outputNames?: string[];
}

export interface ObjectiveNodeData extends BaseNodeData {
  type: 'objective';
  name?: string;
  direction?: 'minimize' | 'maximize';
  weight?: number;
  extractorRef?: string;  // Reference to extractor ID
  outputName?: string;    // Which output from the extractor
  penaltyWeight?: number; // For hard constraints (penalty method)
}

export interface ConstraintNodeData extends BaseNodeData {
  type: 'constraint';
  name?: string;
  operator?: '<' | '<=' | '>' | '>=' | '==';
  value?: number;
}

export interface AlgorithmNodeData extends BaseNodeData {
  type: 'algorithm';
  method?: 'TPE' | 'CMA-ES' | 'NSGA-II' | 'GP-BO' | 'RandomSearch';
  maxTrials?: number;
  // CMA-ES specific
  sigma0?: number;
  restartStrategy?: 'none' | 'ipop' | 'bipop';
  // Weight settings for multi-objective
  objectiveWeights?: Record<string, number>;
}

export interface SurrogateNodeData extends BaseNodeData {
  type: 'surrogate';
  enabled?: boolean;
  modelType?: 'MLP' | 'GNN' | 'Ensemble';
  minTrials?: number;
}

export type CanvasNodeData =
  | ModelNodeData
  | SolverNodeData
  | DesignVarNodeData
  | ExtractorNodeData
  | ObjectiveNodeData
  | ConstraintNodeData
  | AlgorithmNodeData
  | SurrogateNodeData;

export interface CanvasEdge {
  id: string;
  source: string;
  target: string;
  sourceHandle?: string;
  targetHandle?: string;
}

// Valid connections - defines what a node can connect TO (as source)
// Flow: DesignVar -> Model -> Solver -> Extractor -> Objective/Constraint -> Algorithm -> Surrogate
export const VALID_CONNECTIONS: Record<NodeType, NodeType[]> = {
  model: ['solver'],           // Model outputs to Solver
  solver: ['extractor'],       // Solver outputs to Extractor
  designVar: ['model'],        // DesignVar outputs to Model (expressions feed into model)
  extractor: ['objective', 'constraint'], // Extractor outputs to Objective/Constraint
  objective: ['algorithm'],    // Objective outputs to Algorithm
  constraint: ['algorithm'],   // Constraint outputs to Algorithm
  algorithm: ['surrogate'],    // Algorithm outputs to Surrogate
  surrogate: [],               // Surrogate is terminal
};

// Node handle configuration for proper flow direction
export interface HandleConfig {
  id: string;
  label?: string;
}

export interface NodeHandleConfig {
  inputs: HandleConfig[];
  outputs: HandleConfig[];
}

// Define handles for each node type
// Flow: DesignVar(s) -> Model -> Solver -> Extractor(s) -> Objective(s) -> Algorithm
export const NODE_HANDLES: Record<NodeType, NodeHandleConfig> = {
  model: {
    inputs: [{ id: 'params', label: 'Parameters' }],  // Receives from DesignVars
    outputs: [{ id: 'sim', label: 'Simulation' }],    // Sends to Solver
  },
  solver: {
    inputs: [{ id: 'model', label: 'Model' }],        // Receives from Model
    outputs: [{ id: 'results', label: 'Results' }],   // Sends to Extractors
  },
  designVar: {
    inputs: [],                                        // No inputs - this is a source
    outputs: [{ id: 'value', label: 'Value' }],       // Sends to Model
  },
  extractor: {
    inputs: [{ id: 'results', label: 'Results' }],    // Receives from Solver
    outputs: [{ id: 'value', label: 'Value' }],       // Sends to Objective/Constraint
  },
  objective: {
    inputs: [{ id: 'value', label: 'Value' }],        // Receives from Extractor
    outputs: [{ id: 'objective', label: 'Objective' }], // Sends to Algorithm
  },
  constraint: {
    inputs: [{ id: 'value', label: 'Value' }],        // Receives from Extractor
    outputs: [{ id: 'constraint', label: 'Constraint' }], // Sends to Algorithm
  },
  algorithm: {
    inputs: [{ id: 'objectives', label: 'Objectives' }], // Receives from Objectives/Constraints
    outputs: [{ id: 'algo', label: 'Algorithm' }],    // Sends to Surrogate
  },
  surrogate: {
    inputs: [{ id: 'algo', label: 'Algorithm' }],     // Receives from Algorithm
    outputs: [],                                       // No outputs - this is a sink
  },
};