smartai2dian0/src/modules/AISearch/MindMapping.vue

<template>
  <div ref="container" class="mindmap-container"></div>
</template>

<script setup>
import MindMappingData from './MindMappingData.json';
import { onMounted, ref } from 'vue';
import { Rect, Text } from '@antv/g';
import {
  Badge,
  BaseBehavior,
  BaseNode,
  BaseTransform,
  CommonEvent,
  CubicHorizontal,
  ExtensionCategory,
  Graph,
  GraphEvent,
  idOf,
  NodeEvent,
  positionOf,
  register,
  treeToGraphData
} from '@antv/g6';

const container = ref(null);

// 定义样式
const RootNodeStyle = {
  fill: '#EFF0F0',
  labelFill: '#262626',
  labelFontSize: 24,
  labelFontWeight: 600,
  labelOffsetY: 8,
  labelPlacement: 'center',
  ports: [{ placement: 'right' }, { placement: 'left' }],
  radius: 8
};

const NodeStyle = {
  fill: 'transparent',
  labelPlacement: 'center',
  labelFontSize: 16,
  ports: [{ placement: 'right-bottom' }, { placement: 'left-bottom' }]
};

const TreeEvent = {
  COLLAPSE_EXPAND: 'collapse-expand',
  ADD_CHILD: 'add-child'
};

let textShape;

const measureText = (text) => {
  if (!textShape) textShape = new Text({ style: text });
  textShape.attr(text);
  return textShape.getBBox().width;
};

const getNodeWidth = (nodeId, isRoot) => {
  const padding = isRoot ? 40 : 30;
  const nodeStyle = isRoot ? RootNodeStyle : NodeStyle;
  return measureText({ text: nodeId, fontSize: nodeStyle.labelFontSize, fontFamily: 'Gill Sans' }) + padding;
};

const getNodeSize = (nodeId, isRoot) => {
  const width = getNodeWidth(nodeId, isRoot);
  const height = isRoot ? 48 : 32;
  return [width, height];
};

// 自定义节点类
class MindmapNode extends BaseNode {
  static defaultStyleProps = {
    showIcon: false
  };

  constructor(options) {
    Object.assign(options.style, MindmapNode.defaultStyleProps);
    super(options);
  }

  get childrenData() {
    return this.context.model.getChildrenData(this.id);
  }

  get rootId() {
    return idOf(this.context.model.getRootsData()[0]);
  }

  isShowCollapse(attributes) {
    const { collapsed, showIcon } = attributes;
    return !collapsed && showIcon && this.childrenData.length > 0;
  }

  drawCircleShape(attributes, container) {
    const { color = '#000', radius = 10 } = attributes; // 默认颜色和半径
    const [width, height] = this.getSize(attributes);
    const circleStyle = {
      fill: color,
      stroke: '#fff',
      lineWidth: 1,
      r: radius,
      x: width / 2, // 圆心水平位置
      y: height / 2 // 圆心垂直位置
    };
    this.upsert('circle', 'circle', circleStyle, container);
  }

  getCollapseStyle(attributes) {
    const { showIcon, color, direction } = attributes;
    if (!this.isShowCollapse(attributes)) return false;
    const [width, height] = this.getSize(attributes);

    return {
      backgroundFill: color,
      backgroundHeight: 12,
      backgroundWidth: 12,
      cursor: 'pointer',
      fill: '#fff',
      fontFamily: 'iconfont',
      fontSize: 8,
      text: '\ue6e4',
      textAlign: 'center',
      transform: direction === 'left' ? [['rotate', 90]] : [['rotate', -90]],
      visibility: showIcon ? 'visible' : 'hidden',
      x: direction === 'left' ? -6 : width + 6,
      y: height
    };
  }

  drawCollapseShape(attributes, container) {
    const iconStyle = this.getCollapseStyle(attributes);
    const btn = this.upsert('collapse-expand', Badge, iconStyle, container);

    this.forwardEvent(btn, CommonEvent.CLICK, (event) => {
      event.stopPropagation();
      this.context.graph.emit(TreeEvent.COLLAPSE_EXPAND, {
        id: this.id,
        collapsed: !attributes.collapsed
      });
    });
  }

  getCountStyle(attributes) {
    const { collapsed, color, direction } = attributes;
    const count = this.context.model.getDescendantsData(this.id).length;
    if (!collapsed || count === 0) return false;
    const [width, height] = this.getSize(attributes);
    return {
      backgroundFill: color,
      backgroundHeight: 12,
      backgroundWidth: 12,
      cursor: 'pointer',
      fill: '#fff',
      fontSize: 8,
      text: count.toString(),
      textAlign: 'center',
      x: direction === 'left' ? -6 : width + 6,
      y: height
    };
  }

  drawCountShape(attributes, container) {
    const countStyle = this.getCountStyle(attributes);
    const btn = this.upsert('count', Badge, countStyle, container);

    this.forwardEvent(btn, CommonEvent.CLICK, (event) => {
      event.stopPropagation();
      this.context.graph.emit(TreeEvent.COLLAPSE_EXPAND, {
        id: this.id,
        collapsed: false
      });
    });
  }

  getAddStyle(attributes) {
    const { collapsed, showIcon, direction } = attributes;
    if (collapsed || !showIcon) return false;
    const [width, height] = this.getSize(attributes);
    const color = '#ddd';

    const offsetX = this.isShowCollapse(attributes) ? 24 : 12;
    const isRoot = this.id === this.rootId;

    return {
      backgroundFill: '#fff',
      backgroundHeight: 12,
      backgroundLineWidth: 1,
      backgroundStroke: color,
      backgroundWidth: 12,
      cursor: 'pointer',
      fill: color,
      fontFamily: 'iconfont',
      fontSize: 8,
      text: '\ue664',
      textAlign: 'center',
      x: isRoot ? width + 12 : direction === 'left' ? -offsetX : width + offsetX,
      y: isRoot ? height / 2 : height
    };
  }

  getAddBarStyle(attributes) {
    const { collapsed, showIcon, direction, color = '#1783FF' } = attributes;
    if (collapsed || !showIcon) return false;
    const [width, height] = this.getSize(attributes);

    const offsetX = this.isShowCollapse(attributes) ? 12 : 0;
    const isRoot = this.id === this.rootId;

    const HEIGHT = 2;
    const WIDTH = 6;

    return {
      cursor: 'pointer',
      fill:
        direction === 'left'
          ? `linear-gradient(180deg, #fff 20%, ${color})`
          : `linear-gradient(0deg, #fff 20%, ${color})`,
      height: HEIGHT,
      width: WIDTH,
      x: isRoot ? width : direction === 'left' ? -offsetX - WIDTH : width + offsetX,
      y: isRoot ? height / 2 - HEIGHT / 2 : height - HEIGHT / 2,
      zIndex: -1
    };
  }

  drawAddShape(attributes, container) {
    const addStyle = this.getAddStyle(attributes);
    const addBarStyle = this.getAddBarStyle(attributes);
    this.upsert('add-bar', Rect, addBarStyle, container);
    const btn = this.upsert('add', Badge, addStyle, container);

    this.forwardEvent(btn, CommonEvent.CLICK, (event) => {
      event.stopPropagation();
      this.context.graph.emit(TreeEvent.ADD_CHILD, { id: this.id, direction: attributes.direction });
    });
  }

  forwardEvent(target, type, listener) {
    if (target && !Reflect.has(target, '__bind__')) {
      Reflect.set(target, '__bind__', true);
      target.addEventListener(type, listener);
    }
  }

  getKeyStyle(attributes) {
    const [width, height] = this.getSize(attributes);
    const keyShape = super.getKeyStyle(attributes);
    return { width, height, ...keyShape };
  }

  drawKeyShape(attributes, container) {
    const keyStyle = this.getKeyStyle(attributes);
    return this.upsert('key', Rect, keyStyle, container);
  }

  render(attributes = this.parsedAttributes, container = this) {
    super.render(attributes, container);

    this.drawCollapseShape(attributes, container);
    this.drawAddShape(attributes, container);

    this.drawCountShape(attributes, container);
  }
}

// 自定义边类
class MindmapEdge extends CubicHorizontal {
  get rootId() {
    return idOf(this.context.model.getRootsData()[0]);
  }

  getKeyPath(attributes) {
    const path = super.getKeyPath(attributes);
    const isRoot = this.targetNode.id === this.rootId;
    const labelWidth = getNodeWidth(this.targetNode.id, isRoot);

    const [, tp] = this.getEndpoints(attributes);
    const sign = this.sourceNode.getCenter()[0] < this.targetNode.getCenter()[0] ? 1 : -1;
    return [...path, ['L', tp[0] + labelWidth * sign, tp[1]]];
  }
}

// 自定义行为类
class CollapseExpandTree extends BaseBehavior {
  constructor(context, options) {
    super(context, options);
    this.bindEvents();
  }

  update(options) {
    this.unbindEvents();
    super.update(options);
    this.bindEvents();
  }

  bindEvents() {
    const { graph } = this.context;

    graph.on(NodeEvent.POINTER_ENTER, this.showIcon);
    graph.on(NodeEvent.POINTER_LEAVE, this.hideIcon);
    graph.on(TreeEvent.COLLAPSE_EXPAND, this.onCollapseExpand);
    graph.on(TreeEvent.ADD_CHILD, this.addChild);
  }

  unbindEvents() {
    const { graph } = this.context;

    graph.off(NodeEvent.POINTER_ENTER, this.showIcon);
    graph.off(NodeEvent.POINTER_LEAVE, this.hideIcon);
    graph.off(TreeEvent.COLLAPSE_EXPAND, this.onCollapseExpand);
    graph.off(TreeEvent.ADD_CHILD, this.addChild);
  }

  status = 'idle';

  showIcon = (event) => {
    this.setIcon(event, true);
  };

  hideIcon = (event) => {
    this.setIcon(event, false);
  };

  setIcon = (event, show) => {
    if (this.status !== 'idle') return;
    const { target } = event;
    const id = target.id;
    const { graph, element } = this.context;
    graph.updateNodeData([{ id, style: { showIcon: show } }]);
    element.draw({ animation: false, silence: true });
  };

  onCollapseExpand = async (event) => {
    this.status = 'busy';
    const { id, collapsed } = event;
    const { graph } = this.context;
    await graph.frontElement(id);
    if (collapsed) await graph.collapseElement(id);
    else await graph.expandElement(id);
    this.status = 'idle';
  };

  addChild = async (event) => {
    this.status = 'busy';
    const {
      onCreateChild = () => {
        const currentTime = new Date(Date.now()).toLocaleString();
        return { id: `新节点 in ${currentTime}` };
      }
    } = this.options;
    const { graph } = this.context;
    const datum = onCreateChild(event.id);
    const parent = graph.getNodeData(event.id);

    graph.addNodeData([datum]);
    graph.addEdgeData([{ source: event.id, target: datum.id }]);
    graph.updateNodeData([
      {
        id: event.id,
        children: [...(parent.children || []), datum.id],
        style: { collapsed: false, showIcon: false }
      }
    ]);
    await graph.render();
    await graph.focusElement(datum.id);
    this.status = 'idle';
  };
}

// 自定义变换类
class AssignColorByBranch extends BaseTransform {
  static defaultOptions = {
    colors: [
      '#80B4FF', // 淡蓝色
      '#FFC8A2', // 淡橙色
      '#E6C8FF', // 淡紫色
      '#A2E6E6', // 淡青色
      '#C8B4FF', // 淡紫色
      '#FFD878', // 淡黄色
      '#A2E680', // 淡绿色
      '#FFA2D8', // 淡粉色
      '#80C8D8', // 淡蓝色
      '#80E6A2', // 淡绿色
      '#FFD878', // 淡橙色
      '#D880D8', // 淡紫色
      '#80C880', // 淡绿色
      '#FFA280', // 淡红色
      '#80E6E6', // 淡青色
      '#FFA2D8', // 淡粉色
      '#C8C8D8', // 淡蓝色
      '#FFC8C8', // 淡红色
      '#FFD878', // 淡黄色
      '#C8FF80' // 淡绿色
    ] // 增加更多颜色以避免重复
  };

  constructor(context, options) {
    super(context, Object.assign({}, AssignColorByBranch.defaultOptions, options));
  }

  beforeDraw(input) {
    const nodes = this.context.model.getNodeData();
    const edges = this.context.model.getEdgeData();

    if (nodes.length === 0) return input;

    let colorIndex = 0;

    // 为每个节点分配一个独立的颜色
    nodes.forEach((node) => {
      node.style ||= {};
      node.style.color = this.options.colors[colorIndex++ % this.options.colors.length];
    });

    // 为每条边分配颜色，颜色与目标节点的颜色一致
    edges.forEach((edge) => {
      const targetNode = nodes.find((node) => node.id === edge.target);
      if (targetNode) {
        edge.style ||= {};
        edge.style.stroke = targetNode.style.color;
      }
    });

    return input;
  }
}

// 注册扩展
register(ExtensionCategory.NODE, 'mindmap', MindmapNode);
register(ExtensionCategory.EDGE, 'mindmap', MindmapEdge);
register(ExtensionCategory.BEHAVIOR, 'collapse-expand-tree', CollapseExpandTree);
register(ExtensionCategory.TRANSFORM, 'assign-color-by-branch', AssignColorByBranch);

// 获取节点方向
const getNodeSide = (nodeData, parentData) => {
  if (!parentData) return 'center';

  const nodePositionX = positionOf(nodeData)[0];
  const parentPositionX = positionOf(parentData)[0];
  return parentPositionX > nodePositionX ? 'left' : 'right';
};

// 初始化图表
onMounted(async () => {
  // const response = await fetch('https://assets.antv.antgroup.com/g6/algorithm-category.json');
  // const data = await response.json();
  const data = MindMappingData;
  const rootId = data.id;
  const graph = new Graph({
    container: container.value,
    width: container.value.clientWidth || 1000, // 画布宽度，默认值800
    height: container.value.clientHeight || 800, // 画布高度，默认值600
    fitView: true, // 自动缩放以适应视图
    fitCenter: true, // 自动居中
    data: treeToGraphData(data),
    node: {
      type: 'mindmap',
      style: function (d) {
        const direction = getNodeSide(d, this.getParentData(idOf(d), 'tree'));
        const isRoot = idOf(d) === rootId;

        return {
          direction,
          labelText: idOf(d),
          size: getNodeSize(idOf(d), isRoot),
          labelFontFamily: 'Gill Sans',
          labelBackground: true,
          labelBackgroundFill: 'transparent',
          labelPadding: direction === 'left' ? [2, 0, 10, 40] : [2, 40, 10, 0],
          ...(isRoot ? RootNodeStyle : NodeStyle)
        };
      }
    },
    edge: {
      type: 'mindmap',
      style: {
        lineWidth: 3,
        stroke: function (data) {
          return this.getNodeData(data.target).style.color || '#99ADD1';
        }
      }
    },
    layout: {
      type: 'mindmap',
      direction: 'H',
      getHeight: () => 80,
      getWidth: (node) => getNodeWidth(node.id, node.id === rootId),
      getVGap: () => 6,
      getHGap: () => 60,
      animation: false,
      // 新增配置项，使图表从左往右发散
      getSide: () => 'right'
    },
    behaviors: ['drag-canvas', 'zoom-canvas', 'collapse-expand-tree'],
    transforms: ['assign-color-by-branch'],
    animation: {
      duration: 200 // 动画时长，单位为毫秒
    }
  });

  graph.once(GraphEvent.AFTER_RENDER, () => {
    graph.fitView();
  });

  graph.render();
});
</script>

<style scoped>
.mindmap-container {
  width: 100%;
  height: 100%;
}
</style>