@macrostrat/feedback-components 1.0.0

package/src/feedback/edit-state.ts

@@ -0,0 +1,311 @@
+import { TreeData } from "./types";
+import { createContext, Dispatch, useContext, useReducer } from "react";
+import update, { Spec } from "immutability-helper";
+import { EntityType } from "../extractions/types";
+
+export enum ViewMode {
+  Tree = "tree",
+  Graph = "graph",
+}
+
+interface TreeState {
+  initialTree: TreeData[];
+  tree: TreeData[];
+  selectedNodes: number[];
+  entityTypesMap: Map<number, EntityType>;
+  selectedEntityType: EntityType;
+  lastInternalId: number;
+  isSelectingEntityType: boolean;
+  viewMode: ViewMode;
+}
+
+type TextRange = {
+  start: number;
+  end: number;
+  text: string;
+};
+
+type TreeAction =
+  | {
+      type: "move-node";
+      payload: { dragIds: number[]; parentId: number; index: number };
+    }
+  | { type: "delete-node"; payload: { ids: number[] } }
+  | { type: "select-node"; payload: { ids: number[] } }
+  | { type: "toggle-node-selected"; payload: { ids: number[] } }
+  | { type: "set-view-mode"; payload: ViewMode }
+  | { type: "create-node"; payload: TextRange }
+  | { type: "select-entity-type"; payload: EntityType }
+  | { type: "toggle-entity-type-selector"; payload?: boolean | null }
+  | { type: "deselect" }
+  | { type: "reset" };
+
+export type TreeDispatch = Dispatch<TreeAction>;
+
+export function useUpdatableTree(
+  initialTree: TreeData[],
+  entityTypes: Map<number, EntityType>
+): [TreeState, TreeDispatch] {
+  // Get the first entity type
+  const type = entityTypes.values().next().value;
+
+  return useReducer(treeReducer, {
+    initialTree,
+    tree: initialTree,
+    selectedNodes: [],
+    entityTypesMap: entityTypes,
+    selectedEntityType: type,
+    lastInternalId: 0,
+    isSelectingEntityType: false,
+    viewMode: ViewMode.Tree,
+  });
+}
+
+export const TreeDispatchContext = createContext<TreeDispatch | null>(null);
+
+export function useTreeDispatch() {
+  const dispatch = useContext(TreeDispatchContext);
+  if (dispatch == null) {
+    throw new Error("No dispatch context available");
+  }
+  return dispatch;
+}
+
+function treeReducer(state: TreeState, action: TreeAction) {
+  console.log(action);
+  switch (action.type) {
+    case "move-node":
+      // For each node in the tree, if the node is in the dragIds, remove it from the tree and collect it
+      const [newTree, removedNodes] = removeNodes(
+        state.tree,
+        action.payload.dragIds
+      );
+
+      let keyPath: (number | "children")[] = [];
+      if (action.payload.parentId) {
+        keyPath = findNode(newTree, action.payload.parentId);
+        keyPath.push("children");
+      }
+
+      // Add removed nodes to the new tree at the correct location
+      let updateSpec = buildNestedSpec(keyPath, {
+        $splice: [[action.payload.index, 0, ...removedNodes]],
+      });
+
+      return { ...state, tree: update(newTree, updateSpec) };
+    case "delete-node":
+      // For each node in the tree, if the node is in the ids, remove it from the tree
+      const [newTree2, _removedNodes] = removeNodes(
+        state.tree,
+        action.payload.ids
+      );
+      // Get children of the removed nodes
+      // If children are not present elsewhere in the tree, insert them
+
+      const children = _removedNodes
+        .flatMap((node) => node.children ?? [])
+        .filter((child) => !nodeIsInTree(newTree2, child.id));
+
+      // Reset the selection
+
+      return {
+        ...state,
+        tree: [...newTree2, ...children],
+        selectedNodes: state.selectedNodes.filter(
+          (id) => !action.payload.ids.includes(id)
+        ),
+      };
+    case "select-node":
+      const { ids } = action.payload;
+      return { ...state, selectedNodes: ids };
+    // otherwise fall through to toggle-node-selected for a single ID
+    case "toggle-node-selected":
+      const nodesToAdd = action.payload.ids.filter(
+        (id) => !state.selectedNodes.includes(id)
+      );
+      const nodesToKeep = state.selectedNodes.filter(
+        (id) => !action.payload.ids.includes(id)
+      );
+      return { ...state, selectedNodes: [...nodesToKeep, ...nodesToAdd] };
+
+    case "create-node":
+      const newId = state.lastInternalId - 1;
+      const { text, start, end } = action.payload;
+      const node: TreeData = {
+        id: newId,
+        name: text,
+        children: [],
+        indices: [start, end],
+        type: state.selectedEntityType,
+      };
+
+      return {
+        ...state,
+        tree: [...state.tree, node],
+        selectedNodes: [newId],
+        lastInternalId: newId,
+      };
+
+    /** Entity type selection */
+    case "toggle-entity-type-selector":
+      return {
+        ...state,
+        isSelectingEntityType: action.payload ?? !state.isSelectingEntityType,
+      };
+    case "select-entity-type": {
+      // For each selected node, update the type
+      let newTree2 = state.tree;
+      for (let id of state.selectedNodes) {
+        const keyPath = findNode(state.tree, id);
+        const nestedSpec = buildNestedSpec(keyPath, {
+          type: { $set: action.payload },
+        });
+        newTree2 = update(newTree2, nestedSpec);
+      }
+
+      return {
+        ...state,
+        tree: newTree2,
+        selectedEntityType: action.payload,
+      };
+    }
+    case "deselect":
+      return { ...state, selectedNodes: [] };
+    case "reset":
+      return {
+        ...state,
+        tree: state.initialTree,
+        selectedNodes: [],
+      };
+    case "set-view-mode":
+      return { ...state, viewMode: action.payload };
+  }
+}
+
+function nodeIsInTree(tree: TreeData[], id: number): boolean {
+  for (let node of tree) {
+    if (node.id == id) {
+      return true;
+    } else if (node.children) {
+      if (nodeIsInTree(node.children, id)) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+function buildNestedSpec(
+  keyPath: (number | "children")[],
+  innerSpec: Spec<any>
+): Spec<TreeData[]> {
+  // Build a nested object from a key path
+
+  let spec = innerSpec;
+  for (let i = keyPath.length - 1; i >= 0; i--) {
+    spec = { [keyPath[i]]: spec };
+  }
+  return spec as any;
+  // Since we don't have a "children" key at the root, we make the top-level spec an array
+}
+
+function findNode(
+  tree: TreeData[],
+  id: number
+): (number | "children")[] | null {
+  // Find the index of the node with the given id in the tree, returning the key path
+  for (let i = 0; i < tree.length; i++) {
+    if (tree[i].id == id) {
+      return [i];
+    } else if (tree[i].children) {
+      let path = findNode(tree[i].children, id);
+      if (path != null) {
+        return [i, "children", ...path];
+      }
+    }
+  }
+  return null;
+}
+
+function removeNodes(
+  tree: TreeData[],
+  ids: number[]
+): [TreeData[], TreeData[]] {
+  /** Remove nodes with the given ids from the tree and return the new tree and the removed nodes */
+  let newTree: TreeData[] = [];
+  let removedNodes: TreeData[] = [];
+
+  for (let node of tree) {
+    if (ids.includes(node.id)) {
+      removedNodes.push(node);
+    } else {
+      // Recurse into children
+      if (node.children) {
+        let [newChildren, removedChildren] = removeNodes(node.children, ids);
+        node = { ...node, children: newChildren };
+        removedNodes.push(...removedChildren);
+      }
+      newTree.push(node);
+    }
+  }
+
+  return [newTree, removedNodes];
+}
+
+export interface EntityOutput {
+  id: number;
+  type: number | null;
+  txt_range: number[][];
+  name: string;
+  match: any | null;
+  reasoning: string | null;
+}
+
+export interface GraphData {
+  nodes: EntityOutput[];
+  edges: { source: number; dest: number }[];
+}
+
+export function treeToGraph(tree: TreeData[]): GraphData {
+  // Convert the tree to a graph
+  let nodes: EntityOutput[] = [];
+  let edges: { source: number; dest: number }[] = [];
+  const nodeMap = new Map<number, TreeData>();
+
+  for (let node of tree) {
+    // If we've already found an instance of this node, we don't need to record
+    // it again
+    if (nodeMap.has(node.id)) {
+      continue;
+    }
+
+    const { indices, id, name } = node;
+
+    const nodeData: EntityOutput = {
+      id,
+      type: node.type.id,
+      name,
+      txt_range: [indices],
+      reasoning: null,
+      match: node.match,
+    };
+
+    nodeMap.set(node.id, node);
+    nodes.push(nodeData);
+
+    if (node.children) {
+      for (let child of node.children) {
+        edges.push({ source: node.id, dest: child.id });
+      }
+
+      // Now process the children
+      const { nodes: childNodes, edges: childEdges } = treeToGraph(
+        node.children
+      );
+      nodes.push(...childNodes);
+      edges.push(...childEdges);
+    }
+  }
+
+  return { nodes, edges };
+}

package/src/feedback/feedback.module.sass

@@ -0,0 +1,37 @@
+.feedback-component
+  position: relative
+  width: 800px
+
+  & > svg
+    width: 800px
+
+.node
+  cursor: pointer
+
+
+.feedback-text
+  margin-bottom: 2em
+
+.entity-panel
+  position: relative
+  max-height: 600px
+
+.control-panel
+  max-width: 15em
+  position: absolute
+  top: 1em
+  right: 1em
+  padding: 0.2em 0.5em
+
+.entity-panel
+  flex: 1
+  min-height: 100px
+  padding: 1em
+  background: var(--panel-secondary-background-color)
+  border-radius: 4px
+  // Inset box shadow
+  box-shadow: 0 0 0 1px var(--panel-border-color) inset
+
+.selection-tree
+  margin: -1em 0
+  padding: 1em 0

package/src/feedback/graph.ts

@@ -0,0 +1,98 @@
+import { TreeData } from "./types";
+import { treeToGraph } from "./edit-state";
+import h from "@macrostrat/hyper";
+
+import {
+  forceSimulation,
+  SimulationNodeDatum,
+  SimulationLinkDatum,
+  forceCenter,
+  forceLink,
+  forceManyBody,
+  forceCollide,
+} from "d3-force";
+import { useEffect, useState } from "react";
+import { Spinner } from "@blueprintjs/core";
+
+export function GraphView(props: {
+  tree: TreeData[];
+  width: number;
+  height: number;
+}) {
+  // A graph view with react-flow
+  // Get positions of nodes using force simulation
+  const { tree, width, height } = props;
+
+  const [nodes, setNodes] = useState<SimulationNodeDatum[]>(null);
+  const [links, setLinks] = useState<SimulationLinkDatum[]>(null);
+
+  useEffect(() => {
+    const { nodes, edges } = treeToGraph(tree);
+
+    const nodesMap = new Map<number, SimulationNodeDatum>(
+      nodes.map((d) => [d.id, d])
+    );
+
+    const links = edges.map((d) => {
+      return {
+        source: nodesMap.get(d.source),
+        target: nodesMap.get(d.dest),
+        strength: 1,
+      };
+    });
+
+    const simulation = forceSimulation(nodes)
+      .force("link", forceLink(links))
+      .force("charge", forceManyBody().strength(-50))
+      .force("center", forceCenter(width / 2, height / 2))
+      .force("collide", forceCollide().radius(20))
+      .on("tick", () => {
+        // Update the positions of the nodes
+        // setNodes(nodes);
+        console.log("Simulation tick");
+      })
+      .on("end", () => {
+        // Update the positions of the nodes
+        setNodes(nodes);
+        setLinks(links);
+      });
+
+    return () => {
+      simulation.stop();
+    };
+  }, [tree, width, height]);
+
+  if (nodes == null || links == null) {
+    return h(Spinner);
+  }
+
+  console.log("Graph", nodes, links);
+
+  return h("div.graph-view", { style: { width, height } }, [
+    h("svg", { width, height }, [
+      h(
+        "g.nodes",
+        nodes.map((d) => {
+          return h("circle", {
+            cx: d.x,
+            cy: d.y,
+            r: 5,
+            fill: "blue",
+          });
+        })
+      ),
+      h(
+        "g.links",
+        links.map((d) => {
+          return h("line", {
+            x1: d.source.x,
+            y1: d.source.y,
+            x2: d.target.x,
+            y2: d.target.y,
+            stroke: "black",
+          });
+        })
+      ),
+    ]),
+  ]);
+}

package/src/feedback/index.ts

@@ -0,0 +1,271 @@
+import styles from "./feedback.module.sass";
+import hyper from "@macrostrat/hyper";
+
+import { Tree, TreeApi } from "react-arborist";
+import Node from "./node";
+import { FeedbackText } from "./text-visualizer";
+import type { InternalEntity, TreeData } from "./types";
+import type { Entity } from "../extractions";
+import { ModelInfo } from "../extractions";
+import {
+  TreeDispatchContext,
+  treeToGraph,
+  useUpdatableTree,
+  ViewMode,
+} from "./edit-state";
+import { useCallback, useEffect, useRef } from "react";
+import { ButtonGroup, Card, SegmentedControl } from "@blueprintjs/core";
+import { OmniboxSelector } from "./type-selector";
+import {
+  CancelButton,
+  DataField,
+  FlexBox,
+  FlexRow,
+  SaveButton,
+} from "@macrostrat/ui-components";
+import useElementDimensions from "use-element-dimensions";
+import { GraphView } from "./graph";
+
+export type { GraphData } from "./edit-state";
+export { treeToGraph } from "./edit-state";
+export type { TreeData } from "./types";
+
+const h = hyper.styled(styles);
+
+function setsAreTheSame<T>(a: Set<T>, b: Set<T>) {
+  if (a.size !== b.size) return false;
+  for (const item of a) {
+    if (!b.has(item)) return false;
+  }
+  return true;
+}
+
+export function FeedbackComponent({
+  entities = [],
+  text,
+  model,
+  entityTypes,
+  matchComponent,
+  onSave,
+}) {
+  // Get the input arguments
+
+  const [state, dispatch] = useUpdatableTree(
+    entities.map(processEntity) as any,
+    entityTypes
+  );
+
+  const { selectedNodes, tree, selectedEntityType, isSelectingEntityType } =
+    state;
+
+  const [{ width, height }, ref] = useElementDimensions();
+
+  return h(TreeDispatchContext.Provider, { value: dispatch }, [
+    h(FeedbackText, {
+      text,
+      dispatch,
+      // @ts-ignore
+      nodes: tree,
+      selectedNodes,
+    }),
+    h(FlexRow, { alignItems: "baseline", justifyContent: "space-between" }, [
+      h(ModelInfo, { data: model }),
+      h(SegmentedControl, {
+        options: [
+          { label: "Tree", value: "tree" },
+          { label: "Graph", value: "graph" },
+        ],
+        value: state.viewMode,
+        small: true,
+        onValueChange(value: ViewMode) {
+          console.log("Setting view mode", value);
+          dispatch({ type: "set-view-mode", payload: value });
+        },
+      }),
+    ]),
+    h(
+      "div.entity-panel",
+      {
+        ref,
+      },
+      [
+        h(Card, { className: "control-panel" }, [
+          h(
+            ButtonGroup,
+            {
+              vertical: true,
+              fill: true,
+              minimal: true,
+              alignText: "left",
+            },
+            [
+              h(
+                CancelButton,
+                {
+                  icon: "trash",
+                  disabled: state.initialTree == state.tree,
+                  onClick() {
+                    dispatch({ type: "reset" });
+                  },
+                },
+                "Reset"
+              ),
+              h(
+                SaveButton,
+                {
+                  onClick() {
+                    onSave(state.tree);
+                  },
+                  disabled: state.initialTree == state.tree,
+                },
+                "Save"
+              ),
+            ]
+          ),
+          h(EntityTypeSelector, {
+            entityTypes,
+            selected: selectedEntityType,
+            onChange(payload) {
+              dispatch({ type: "select-entity-type", payload });
+            },
+            isOpen: isSelectingEntityType,
+            setOpen: (isOpen: boolean) =>
+              dispatch({
+                type: "toggle-entity-type-selector",
+                payload: isOpen,
+              }),
+          }),
+        ]),
+        h.if(state.viewMode == "tree")(ManagedSelectionTree, {
+          selectedNodes,
+          dispatch,
+          tree,
+          width,
+          height,
+          matchComponent,
+        }),
+        h.if(state.viewMode == "graph")(GraphView, {
+          tree,
+          width,
+          height,
+        }),
+      ]
+    ),
+  ]);
+}
+
+function processEntity(entity: Entity): InternalEntity {
+  // @ts-ignore
+  return {
+    ...entity,
+    // @ts-ignore
+    term_type: entity.type.name,
+    txt_range: [entity.indices],
+    children: entity.children?.map(processEntity) ?? [],
+  };
+}
+
+function EntityTypeSelector({
+  entityTypes,
+  selected,
+  isOpen,
+  setOpen,
+  onChange,
+}) {
+  // Show all entity types when selected is null
+  const _selected = selected != null ? selected : undefined;
+  return h(DataField, { label: "Entity type", inline: true }, [
+    h(
+      "code.bp5-code",
+      {
+        onClick() {
+          setOpen((d) => !d);
+        },
+      },
+      selected.name
+    ),
+    h(OmniboxSelector, {
+      isOpen,
+      items: Array.from(entityTypes.values()),
+      selectedItem: _selected,
+      onSelectItem(item) {
+        setOpen(false);
+        onChange(item);
+      },
+      onClose() {
+        setOpen(false);
+      },
+    }),
+  ]);
+}
+
+function ManagedSelectionTree(props) {
+  const {
+    selectedNodes,
+    dispatch,
+    tree,
+    height,
+    width,
+    matchComponent,
+    ...rest
+  } = props;
+
+  const ref = useRef<TreeApi<TreeData>>();
+
+  const _Node = useCallback(
+    (props) => h(Node, { ...props, matchComponent }),
+    [matchComponent]
+  );
+
+  useEffect(() => {
+    if (ref.current == null) return;
+    // Check if selection matches current
+    const selection = new Set(selectedNodes.map((d) => d.toString()));
+    const currentSelection = ref.current.selectedIds;
+    if (setsAreTheSame(selection, currentSelection)) return;
+    // If the selection is the same, do nothing
+
+    // Set selection
+    ref.current.setSelection({
+      ids: selectedNodes.map((d) => d.toString()),
+      anchor: null,
+      mostRecent: null,
+    });
+  }, [selectedNodes]);
+
+  return h(Tree, {
+    className: "selection-tree",
+    height,
+    width,
+    ref,
+    data: tree,
+    onMove({ dragIds, parentId, index }) {
+      dispatch({
+        type: "move-node",
+        payload: {
+          dragIds: dragIds.map((d) => parseInt(d)),
+          parentId: parentId ? parseInt(parentId) : null,
+          index,
+        },
+      });
+    },
+    onDelete({ ids }) {
+      dispatch({
+        type: "delete-node",
+        payload: { ids: ids.map((d) => parseInt(d)) },
+      });
+    },
+    onSelect(nodes) {
+      let ids = nodes.map((d) => parseInt(d.id));
+      if (ids.length == 1 && ids[0] == selectedNodes[0]) {
+        // Deselect
+        ids = [];
+      }
+      dispatch({ type: "select-node", payload: { ids } });
+    },
+    children: _Node,
+    idAccessor(d: TreeData) {
+      return d.id.toString();
+    },
+  });
+}