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

import { useEffectWhen } from '@/hooks';
import { ResourceVersionData } from '@/pages/Dataset/config';
import { getModelAtlasReq } from '@/services/dataset/index.js';
import themes from '@/styles/theme.less';
import { to } from '@/utils/promise';
import G6, { G6GraphEvent, Graph } from '@antv/g6';
// @ts-ignore
import { ResourceInfoTabKeys } from '@/pages/Dataset/components/ResourceIntro';
import { Flex, Select } from 'antd';
import { useEffect, useRef, useState } from 'react';
import GraphLegend from '../GraphLegend';
import NodeTooltips from '../NodeTooltips';
import styles from './index.less';
import type { ModelDepsData, ProjectDependency, TrainDataset } from './utils';
import { NodeType, getGraphData, nodeHeight, nodeWidth, normalizeTreeData } from './utils';

type modeModelEvolutionProps = {
  resourceId: number;
  versionList: ResourceVersionData[];
  version?: string;
  isActive: boolean;
  onVersionChange: (version: string) => void;
};

let graph: Graph;
function ModelEvolution({
  resourceId,
  versionList,
  version,
  isActive,
  onVersionChange,
}: modeModelEvolutionProps) {
  const graphRef = useRef<HTMLDivElement>(null);
  const [showNodeTooltip, setShowNodeTooltip] = useState(false);
  const [enterTooltip, setEnterTooltip] = useState(false);
  const [nodeTooltipX, setNodeToolTipX] = useState(0);
  const [nodeTooltipY, setNodeToolTipY] = useState(0);
  const [hoverNodeData, setHoverNodeData] = useState<
    ModelDepsData | ProjectDependency | TrainDataset | undefined
  >(undefined);

  useEffect(() => {
    initGraph();
    const changeSize = () => {
      if (!graph || graph.get('destroyed')) return;
      if (!graphRef.current) return;
      graph.changeSize(graphRef.current.clientWidth, graphRef.current.clientHeight);
      graph.fitView();
    };

    window.addEventListener('resize', changeSize);
    return () => {
      window.removeEventListener('resize', changeSize);
    };
  }, []);

  useEffectWhen(
    () => {
      if (version) {
        getModelAtlas();
      } else {
        clearGraphData();
      }
    },
    [resourceId, version],
    isActive,
  );

  // 初始化图
  const initGraph = () => {
    graph = new G6.Graph({
      container: graphRef.current!,
      width: graphRef.current!.clientWidth,
      height: graphRef.current!.clientHeight,
      fitView: true,
      fitViewPadding: [50, 100, 50, 100],
      minZoom: 0.5,
      maxZoom: 5,
      defaultNode: {
        type: 'rect',
        size: [nodeWidth, nodeHeight],
        anchorPoints: [
          [0, 0.5],
          [1, 0.5],
          [0.5, 0],
          [0.5, 1],
        ],
        style: {
          fill: themes['primaryColor'],
          lineWidth: 0,
          radius: 6,
          cursor: 'pointer',
        },
        labelCfg: {
          position: 'center',
          style: {
            fill: '#ffffff',
            fontSize: 8,
            textAlign: 'center',
            cursor: 'pointer',
          },
        },
      },
      defaultEdge: {
        type: 'cubic-horizontal',
        labelCfg: {
          autoRotate: true,
        },
        style: {
          stroke: '#a2c1ff',
          lineWidth: 1,
        },
      },
      modes: {
        default: [
          'drag-canvas',
          'zoom-canvas',
          // {
          //   type: 'collapse-expand',
          //   onChange(item?: Item, collapsed?: boolean) {
          //     const data = item!.getModel();
          //     data.collapsed = collapsed;
          //     return true;
          //   },
          // },
        ],
      },
    });

    bindEvents();
  };

  // 绑定事件
  const bindEvents = () => {
    graph.on('node:mouseenter', (e: G6GraphEvent) => {
      const nodeItem = e.item;
      graph.setItemState(nodeItem, 'hover', true);

      const model = nodeItem.getModel() as ModelDepsData;
      const { x, y } = model;
      const point = graph.getCanvasByPoint(x!, y!);
      const zoom = graph.getZoom();
      // 更加缩放，调整 tooltip 位置
      const offsetX = (nodeWidth * zoom) / 4;
      const offsetY = (nodeHeight * zoom) / 4;

      const canvasWidth = graphRef.current!.clientWidth;
      if (point.x + 300 > canvasWidth) {
        point.x = canvasWidth - 300;
      }

      setHoverNodeData(model);
      setNodeToolTipX(point.x + offsetX);
      setNodeToolTipY(graphRef.current!.clientHeight - point.y + offsetY);
      setShowNodeTooltip(true);
    });

    graph.on('node:mouseleave', (e: G6GraphEvent) => {
      const nodeItem = e.item;
      graph.setItemState(nodeItem, 'hover', false);
      setShowNodeTooltip(false);
    });

    graph.on('node:click', (e: G6GraphEvent) => {
      const nodeItem = e.item;
      const model = nodeItem.getModel();
      const { model_type } = model;
      const { origin } = location;
      let url: string = '';
      switch (model_type) {
        case NodeType.children:
        case NodeType.parent: {
          const { current_model_id, version } = model as ModelDepsData;
          url = `${origin}/dataset/model/${current_model_id}?tab=${ResourceInfoTabKeys.Evolution}&version=${version}`;
          break;
        }
        case NodeType.project: {
          const { url: projectUrl } = model as ProjectDependency;
          url = projectUrl;
          break;
        }
        case NodeType.trainDataset:
        case NodeType.testDataset: {
          const { dataset_id, dataset_version } = model as TrainDataset;
          url = `${origin}/dataset/dataset/${dataset_id}?tab=${ResourceInfoTabKeys.Version}&version=${dataset_version}`;
          break;
        }
        default:
          break;
      }

      if (url) {
        window.open(url, '_blank');
      }
    });

    // 鼠标滚轮缩放时，隐藏 tooltip
    graph.on('wheelzoom', () => {
      setShowNodeTooltip(false);
      setEnterTooltip(false);
    });
  };

  const handleTooltipsMouseEnter = () => {
    setEnterTooltip(true);
  };

  const handleTooltipsMouseLeave = () => {
    setEnterTooltip(false);
  };

  // 获取模型依赖
  const getModelAtlas = async () => {
    const params = {
      current_model_id: resourceId,
      version,
    };
    const [res] = await to(getModelAtlasReq(params));
    if (res && res.data) {
      const data = normalizeTreeData(res.data);
      const graphData = getGraphData(data);

      graph.data(graphData);
      graph.render();
      graph.fitView();
    } else {
      clearGraphData();
    }
  };

  // 请求失败或者版本不存在时，清除图形
  function clearGraphData() {
    graph.data({
      nodes: [],
      edges: [],
    });
    graph.render();
    graph.fitView();
  }

  return (
    <div className={styles['model-evolution']}>
      <Flex align="center" className={styles['model-evolution__top']}>
        <span style={{ marginRight: '10px' }}>版本号：</span>
        <Select
          placeholder="请选择版本号"
          style={{ width: '160px', marginRight: '20px' }}
          value={version}
          allowClear
          onChange={onVersionChange}
          options={versionList}
        />
        <GraphLegend style={{ marginRight: 0, marginLeft: 'auto' }}></GraphLegend>
      </Flex>
      <div className={styles['model-evolution__graph']} id="canvas" ref={graphRef}></div>
      {(showNodeTooltip || enterTooltip) && (
        <NodeTooltips
          x={nodeTooltipX}
          y={nodeTooltipY}
          data={hoverNodeData!}
          onMouseEnter={handleTooltipsMouseEnter}
          onMouseLeave={handleTooltipsMouseLeave}
        />
      )}
    </div>
  );
}

export default ModelEvolution;