path: root/src/zoom/StanzaClause.tsx

import React, { memo } from "react";
import { motion } from "motion/react";
import "./parsing.css";
import type { ViewProps } from "./logic/types";
import Word from "./Word";
import { clauseVariants, createHoverEffect } from "./animations";
import { useZoom } from "./hooks/useZoom";
import { NLP } from "sortug-ai";

// Function to check if a string is a punctuation character

const isLeaf = (node: NLP.Stanza.TreeNode): boolean => {
  return node.children && node.children.length === 0;
};
const toIgnore = ["root", "s"];
// Component to render each node in the constituency tree
const TreeNode = ({
  node,
  nest = 0,
  idx,
}: {
  idx: number;
  nest?: number;
  node: NLP.Stanza.TreeNode;
}) => {
  const neatChildren = node.children.reduce(
    (acc: NLP.Stanza.TreeNode[], item) => {
      if (NLP.isPunctuation(item.label)) return acc;
      else return [...acc, item];
    },
    [],
  );

  return !isLeaf(node) ? (
    <BranchNode
      node={{ ...node, children: neatChildren }}
      nest={nest}
      idx={idx}
    />
  ) : (
    <LeafNode text={node.label} />
  );
};

const BranchNode = ({
  node,
  nest = 0,
  idx,
}: {
  idx: number;
  nest?: number;
  node: NLP.Stanza.TreeNode;
}) => {
  const { viewState, handleElementClick } = useZoom();
  const { level, cIndex } = viewState;
  const selected = cIndex === idx;
  const isFocused = level === "clause" && selected;
  const color = `rgba(${SYNTAX_COLORS[node.label]})` || "100, 100, 100";
  const style: any = { "--clause-underline-color": color };
  const out = toIgnore.includes(node.label.toLowerCase());
  const isPunct = NLP.isPunctuation(node.label);
  if (isPunct) return null;
  const clauseLabel = out ? null : NLP.Stanza.oneDescendant(node) ? (
    <div className="word-pos">{node.label}</div>
  ) : (
    <div className="clause-pos">{node.label}</div>
  );
  return (
    <motion.div
      style={style}
      className={`clause${selected ? " selected" : ""}`}
      custom={selected}
      variants={clauseVariants}
      initial="sentence"
      animate={level}
      onClick={(e) => handleElementClick(e, idx)}
      whileHover={createHoverEffect(
        level,
        "sentence",
        SYNTAX_COLORS[node.label],
      )}
    >
      <div className="clause-inner">
        <div style={{ display: "flex" }}>
          {node.children.map((child, childIdx) => (
            <TreeNode
              key={childIdx}
              node={child}
              nest={nest + 1}
              idx={childIdx + idx + 1} // Ensure unique index
            />
          ))}
        </div>
        {clauseLabel}
      </div>
    </motion.div>
  );
};
function LeafNode({ text }: { text: string }) {
  return <motion.div className="tree-leaf">{text}</motion.div>;
}

// Main Clause component
interface Props extends ViewProps {
  sentence: NLP.Stanza.Sentence;
  data: NLP.Stanza.TreeNode;
}

function SimpleClause(props: Props) {
  const { sentence, data, context, rawText, idx } = props;
  const { viewState, handleElementClick } = useZoom();
  const { level, cIndex } = viewState;
  const selected = cIndex === idx;
  const isFocused = level === "clause" && selected;
  console.log({ viewState, rawText, f: isFocused, idx, cIndex });

  const words = extractWordsFromTree(sentence, data);
  const segmented = words.map((w) => w.text);
  console.log({ words });

  return (
    <motion.div
      className={`clause-container ${selected ? "selected" : ""}`}
      custom={selected}
      variants={clauseVariants}
      initial="sentence"
      animate={level}
      onClick={(e) => handleElementClick(e, idx)}
      whileHover={createHoverEffect(level, "sentence", "220, 200, 255")}
    >
      {level === "sentence" ? (
        <span className="clause">{rawText}</span>
      ) : (
        <div className="words-container">
          {words.map((word, wordIdx) => (
            <Word
              {...props}
              key={word.text + wordIdx}
              idx={wordIdx}
              rawText={word.text}
              word={word}
              context={{
                idx: wordIdx,
                parentText: rawText,
                segmented,
              }}
            />
          ))}
        </div>
      )}
    </motion.div>
  );
}
function RecursiveClause(props: Props) {
  const { sentence, data, context, rawText, idx } = props;
  const { viewState, handleElementClick } = useZoom();
  const { level, cIndex } = viewState;
  const selected = cIndex === idx;
  const isFocused = level === "clause" && selected;

  // If we're at the word level, display words instead of the constituency tree
  if (level === "word" && selected && data.children) {
    // Extract word objects if available
    const words = extractWordsFromTree(sentence, data);
    if (words.length > 0) {
      return (
        <div className="words-container">
          {words.map((word, wordIdx) => (
            <Word
              {...props}
              key={word.text + wordIdx}
              idx={wordIdx}
              rawText={word.text}
              word={word}
              context={{
                idx: wordIdx,
                parentText: rawText,
                segmented: words.map((w) => w.text),
              }}
            />
          ))}
        </div>
      );
    }
  }

  return (
    <>
      <TreeNode node={data} idx={idx} />
    </>
  );
}

// Helper function to extract words from the tree
function extractWordsFromTree(
  sentence: NLP.Stanza.Sentence,
  node: NLP.Stanza.TreeNode,
): NLP.Stanza.Word[] {
  const words: NLP.Stanza.Word[] = [];

  if (node.label && node.children.length === 0) {
    // Check if this is a punctuation node
    if (!NLP.isPunctuation(node.label)) {
      // Find the matching word in the sentence
      const word = sentence.words.find((w) => w.text === node.label);
      if (word) {
        words.push(word);
      }
    }
  }

  // Recursively process children
  if (node.children) {
    for (const child of node.children) {
      words.push(...extractWordsFromTree(sentence, child));
    }
  }

  // Remove duplicates by word id
  const uniqueWords = Array.from(
    new Map(words.map((word) => [word?.id, word])).values(),
  ).filter(Boolean) as NLP.Stanza.Word[];

  return uniqueWords;
}

// Syntax highlighting colors - reusing from Stanza utils
const SYNTAX_COLORS: any = {
  // Sentence - cornflower blue
  S: "100,149,237",
  // Sentence - cornflower blue
  SBAR: "100,149,237",
  // Sentence - cornflower blue
  SBARQ: "100,149,237",
  // Noun Phrase - coral
  NP: "255,127,80",
  // Verb Phrase - lime green
  VP: "50,205,50",
  // Prepositional Phrase - medium purple
  PP: "147,112,219",
  // Adjective Phrase - gold
  AP: "255,215,0",
  // Adverb Phrase - hot pink
  AVP: "255,105,180",
  // Noun - light salmon
  NN: "255,160,122",
  // Verb - light green
  V: "144,238,144",
  // Verb - light green
  VB: "144,238,144",
  // Verb - light green
  VBP: "144,238,144",
  // Verb - light green
  VBG: "144,238,144",
  // Verb - light green
  VBZ: "144,238,144",
  // Verb - light green
  VBD: "144,238,144",
  // Verb - light green
  VBN: "144,238,144",
  // Adjective - khaki
  JJ: "240,230,140",
  // Adverb - plum
  ADV: "221,160,221",
  // Preposition - light sky blue
  PR: "135,206,250",
  // Preposition - light sky blue
  IN: "135,206,250",
  // Preposition - light sky blue
  TO: "135,206,250",
  // Determiner - light gray
  DT: "211,211,211",
  // Personal pronoun - thistle
  PPN: "216,191,216",
  // Coordinating conjunction - dark gray
  CC: "169,169,169",
};

export default memo(RecursiveClause);