ci4sManagement-cloud/react-ui/src/pages/Model/components/ModelEvolution/utils.tsx

import { changePropertyName, fittingString } from '@/utils';
import { EdgeConfig, GraphData, LayoutConfig, NodeConfig, TreeGraphData, Util } from '@antv/g6';
// @ts-ignore
import Hierarchy from '@antv/hierarchy';

export const nodeWidth = 110;
export const nodeHeight = 50;
export const vGap = nodeHeight + 20;
export const hGap = nodeHeight + 20;
export const ellipseWidth = nodeWidth;

// 数据集节点矩形数组
const datasetRects: Rect[] = [];

export enum NodeType {
  current = 'current',
  parent = 'parent',
  children = 'children',
  project = 'project',
  trainDataset = 'trainDataset',
  testDataset = 'testDataset',
}

export type Rect = {
  x: number;
  y: number;
  width: number;
  height: number;
};

export type TrainTask = {
  ins_id: number;
  name: string;
  task_id: string;
};

export interface TrainDataset extends NodeConfig {
  dataset_id: number;
  dataset_name: string;
  dataset_version: string;
  model_type: NodeType;
}

export interface ProjectDependency extends NodeConfig {
  url: string;
  name: string;
  branch: string;
  model_type: NodeType;
}

export type ModalDetail = {
  name: string;
  available_range: number;
  file_name: string;
  file_size: string;
  description: string;
  model_type_name: string;
  model_tag_name: string;
  create_time: string;
};

export interface ModelDepsAPIData {
  current_model_id: number;
  version: string;
  exp_ins_id: number;
  model_type: NodeType;
  current_model_name: string;
  project_dependency: ProjectDependency;
  test_dataset: TrainDataset[];
  train_dataset: TrainDataset[];
  train_task: TrainTask;
  model_version_dependcy_vo: ModalDetail;
  children_models: ModelDepsAPIData[];
  parent_models: ModelDepsAPIData[];
}

export interface ModelDepsData extends Omit<ModelDepsAPIData, 'children_models'>, TreeGraphData {
  children: ModelDepsData[];
}

// 规范化子数据
export function normalizeChildren(data: ModelDepsData[]) {
  if (Array.isArray(data)) {
    data.forEach((item) => {
      item.model_type = NodeType.children;
      item.id = `$M_${item.current_model_id}_${item.version}`;
      item.label = getLabel(item);
      item.style = getStyle(NodeType.children);
      normalizeChildren(item.children);
    });
  }
}

// 获取 label
export function getLabel(node: ModelDepsData | ModelDepsAPIData) {
  return (
    fittingString(`${node.model_version_dependcy_vo.name ?? ''}`, nodeWidth - 12, 8) +
    '\n' +
    fittingString(`${node.version}`, nodeWidth - 12, 8)
  );
}

// 获取 style
export function getStyle(model_type: NodeType) {
  let fill = '';
  switch (model_type) {
    case NodeType.current:
      fill = 'l(0) 0:#72a1ff 1:#1664ff';
      break;
    case NodeType.parent:
      fill = 'l(0) 0:#93dfd1 1:#43c9b1';
      break;
    case NodeType.children:
      fill = 'l(0) 0:#72b4ff 1:#169aff';
      break;
    case NodeType.project:
      fill = 'l(0) 0:#b3a9ff 1:#8981ff';
      break;
    case NodeType.trainDataset:
      fill = '#a5d878';
      break;
    case NodeType.testDataset:
      fill = '#d8b578';
      break;
    default:
      break;
  }
  return {
    fill,
  };
}

// 将后台返回的数据转换成树形数据
export function normalizeTreeData(
  apiData: ModelDepsAPIData,
  currentNodeName: string,
): ModelDepsData {
  // 将 children_models 转换成 children
  let normalizedData = changePropertyName(apiData, {
    children_models: 'children',
  }) as ModelDepsData;

  // 设置当前模型的数据
  normalizedData.model_type = NodeType.current;
  normalizedData.current_model_name = currentNodeName;
  normalizedData.id = `$M_${normalizedData.current_model_id}_${normalizedData.version}`;
  normalizedData.label = getLabel(normalizedData);
  normalizedData.style = getStyle(NodeType.current);
  // let first1 = { ...normalizedData.children[0] };
  // let first2 = { ...normalizedData.children[0] };
  // let first3 = { ...normalizedData.children[0] };
  // first1.current_model_id = 202020;
  // first2.current_model_id = 202021;
  // first3.current_model_id = 202022;
  // normalizedData.children.push(first1, first2, first3);

  normalizeChildren(normalizedData.children as ModelDepsData[]);

  // 将 parent_models 转换成树形结构
  let parent_models = normalizedData.parent_models || [];
  while (parent_models.length > 0) {
    const parent = parent_models[0];
    normalizedData = {
      ...parent,
      model_type: NodeType.parent,
      id: `$M_${parent.current_model_id}_${parent.version}`,
      label: getLabel(parent),
      style: getStyle(NodeType.parent),
      children: [
        {
          ...normalizedData,
          parent_models: [],
        },
      ],
    };
    parent_models = normalizedData.parent_models || [];
  }
  return normalizedData;
}

// 将树形数据，使用 Hierarchy 进行布局，计算出坐标，然后转换成 G6 的数据
export function getGraphData(data: ModelDepsData): GraphData {
  const config = {
    direction: 'LR',
    getHeight: () => nodeHeight,
    getWidth: () => nodeWidth,
    getVGap: () => vGap / 2,
    getHGap: () => hGap / 2,
  };

  // 树形布局计算出坐标
  const treeLayoutData: LayoutConfig = Hierarchy['compactBox'](data, config);

  const nodes: NodeConfig[] = [];
  const edges: EdgeConfig[] = [];
  Util.traverseTree(treeLayoutData, (node: NodeConfig, parent: NodeConfig) => {
    const data = node.data as ModelDepsData;
    console.log('data', data);

    // 当前模型显示数据集和项目
    if (data.model_type === NodeType.current) {
      addDatasetDependency(data, node, nodes, edges);
      addProjectDependency(data, node, nodes, edges);
    } else if (data.model_type === NodeType.children) {
      adjustChildrenPosition(node);
    }
    nodes.push({
      ...data,
      x: node.x,
      y: node.y,
    });
    if (parent) {
      edges.push({
        source: parent.id,
        target: node.id,
      });
    }
  });
  return { nodes, edges };
}

// 将数据集转换成 G6 的数据
const addDatasetDependency = (
  data: ModelDepsData,
  currentNode: NodeConfig,
  nodes: NodeConfig[],
  edges: EdgeConfig[],
) => {
  const { train_dataset, test_dataset } = data;
  train_dataset.forEach((item) => {
    item.id = `$DTrain_${item.dataset_id}`;
    item.model_type = NodeType.trainDataset;
    item.style = getStyle(NodeType.trainDataset);
  });
  test_dataset.forEach((item) => {
    item.id = `$DTest_${item.dataset_id}`;
    item.model_type = NodeType.testDataset;
    item.style = getStyle(NodeType.testDataset);
  });

  datasetRects.length = 0;
  const len = train_dataset.length + test_dataset.length;
  [...train_dataset, ...test_dataset].forEach((item, index) => {
    const node = { ...item };
    node.type = 'ellipse';
    node.size = [ellipseWidth, nodeHeight];
    node.label =
      fittingString(node.dataset_name, ellipseWidth - 12, 8) +
      '\n' +
      fittingString(node.dataset_version, ellipseWidth - 12, 8);

    const half = len / 2 - 0.5;
    node.x = currentNode.x! - (half - index) * (ellipseWidth + hGap);
    node.y = currentNode.y! - nodeHeight - vGap;
    nodes.push(node);
    edges.push({
      source: currentNode.id,
      target: node.id,
      sourceAnchor: 2,
      targetAnchor: 3,
      type: 'cubic-vertical',
    });
    datasetRects.push({
      x: node.x - ellipseWidth / 2,
      y: node.y - nodeHeight / 2,
      width: ellipseWidth,
      height: nodeHeight,
    });
  });
};

// 将模型依赖数据转换成 G6 的数据
const addProjectDependency = (
  data: ModelDepsData,
  currentNode: NodeConfig,
  nodes: NodeConfig[],
  edges: EdgeConfig[],
) => {
  const { project_dependency } = data;
  if (project_dependency?.url) {
    const node = { ...project_dependency };
    node.id = `$P_${node.url}`;
    node.model_type = NodeType.project;
    node.type = 'rect';
    node.label = fittingString(node.name, nodeWidth - 12, 8);
    node.style = getStyle(NodeType.project);
    node.style.radius = nodeHeight / 2;
    node.x = currentNode.x;
    node.y = currentNode.y! + nodeHeight + vGap;

    nodes.push(node);
    edges.push({
      source: currentNode.id,
      target: node.id,
      sourceAnchor: 3,
      targetAnchor: 2,
      type: 'cubic-vertical',
    });
  }
};

// 判断两个矩形是否相交
function isRectanglesIntersect(rect1: Rect, rect2: Rect) {
  return !(
    rect1.x + rect1.width < rect2.x ||
    rect1.x > rect2.x + rect2.width ||
    rect1.y + rect1.height < rect2.y ||
    rect1.y > rect2.y + rect2.height
  );
}

// 判断子节点是否与数据集节点重叠
function isChildrenIntersectDataset(rects: Rect[], childrenRect: Rect) {
  for (const r of rects) {
    if (isRectanglesIntersect(r, childrenRect)) {
      return r;
    }
  }

  return null;
}

// 计算子节点位置
function adjustChildrenPosition(node: NodeConfig) {
  const nodeRect = {
    x: node.x! - nodeWidth / 2,
    y: node.y! - nodeHeight / 2,
    width: nodeWidth,
    height: nodeHeight,
  };
  const overlapRect = isChildrenIntersectDataset(datasetRects, nodeRect);
  if (overlapRect) {
    const offsetY = nodeRect.y - overlapRect.y;
    const space = 10; //(vGap + Math.abs(offsetY) - nodeHeight) / 2;
    if (offsetY >= 0) {
      node.y = node.y! + (nodeHeight - offsetY + space);
    } else {
      node.y = node.y! - (nodeHeight - Math.abs(offsetY) + space);
    }
  }
}