smartai2dian0/src/modules/ai-search/MindMapping.vue

<template>
  <div v-show="!data && loading" class="progress">
    <a-progress type="circle" width="250" stroke-width="10" :percent="progress / 100" :show-text="false" />
    <h1>{{ `${progress}%` }}</h1>
  </div>

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

<script setup>
import MindMappingData from './MindMappingData.json';
import { onMounted, watch, onUnmounted, 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 props = defineProps({
  data: Object,
  loading: Boolean
});

const container = ref(null);
const progress = ref(0);
const interval_ref = ref(null);

const startProgress = () => {
  interval_ref.value = setInterval(() => {
    progress.value += Math.random() > 0.5 ? 3 : 2;
    if (progress.value > 98) {
      progress.value = 98; // 确保 progress 不会超过 98
      clearInterval(interval_ref.value);
    }
  }, 1000);
};

onMounted(() => {
  if (props.loading && !props.data) {
    startProgress();
  }
});

onUnmounted(() => {
  clearInterval(interval_ref.value);
});

watch(
  () => props.loading,
  (new_loading) => {
    if (new_loading && !props.data) {
      startProgress();
    } else {
      clearInterval(interval_ref.value);
    }
  }
);

watch(
  () => props.data,
  (new_data) => {
    if (new_data) {
      initGraph(new_data);
    }
  }
);
// 定义样式
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 text_shape;

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

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

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

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

  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;
  }

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

    // TODO 有边框的小白点，stroke失效
    return {
      backgroundFill: color, // 设置为白色
      backgroundHeight: 12, // 调整为小圆点的大小
      backgroundWidth: 12, // 调整为小圆点的大小
      cursor: 'pointer',
      stroke: color, // 使用传入的颜色作为边框颜色
      lineWidth: 4, // 设置边框宽度
      fontSize: 0, // 不需要字体大小
      text: '', // 移除文本
      textAlign: 'center',
      transform: [], // 移除旋转效果
      visibility: 'visible',
      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' ? -4 : width + 4,
      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;
    const isLeaf = this.childrenData.length === 0; // 检查节点是否为叶子节点
    if (collapsed || !showIcon || !isLeaf) return false; // 仅为叶子节点显示加号按钮
    const [width, height] = this.getSize(attributes);

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

    return {
      backgroundFill: '#1783FF',
      backgroundHeight: 14,
      backgroundLineWidth: 1,
      backgroundStroke: '#1783FF',
      backgroundWidth: 14,
      cursor: 'pointer',
      fill: '#fff',
      fontSize: 14, // 调整字体大小以适应加号
      text: '+', // 使用普通加号
      textAlign: 'center',
      x: isRoot ? width + 12 : direction === 'left' ? -offsetX : width + offsetX,
      y: isRoot ? height / 2 : height
    };
  }

  drawAddShape(attributes, container) {
    const addStyle = this.getAddStyle(attributes);
    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 = (node_data, parent_data) => {
  if (!parent_data) return 'center';

  const node_position_x = positionOf(node_data)[0];
  const parent_position_x = positionOf(parent_data)[0];
  return parent_position_x > node_position_x ? 'left' : 'right';
};

// 初始化图表
const initGraph = (data) => {
  const root_id = 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 is_root = idOf(d) === root_id;

        return {
          direction,
          labelText: idOf(d),
          size: getNodeSize(idOf(d), is_root),
          labelFontFamily: 'Gill Sans',
          labelBackground: true,
          labelBackgroundFill: 'transparent',
          labelPadding: direction === 'left' ? [2, 0, 10, 40] : [2, 40, 10, 0],
          color: d.style.color,
          ...(is_root ? 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 === root_id),
      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%;
}

.progress {
  margin-top: 5%;
  position: relative;
  text-align: center;
}

.progress h1 {
  position: absolute;
  top: 50%;
  left: 50%;
  transform: translate(-50%, -50%);
  margin: 0;
}
</style>