SEECBC/FrontEnd/components/playground/gradient-descent/step-chart.ts

import * as d3 from 'd3'

type NodeType = 'normal' | 'start' | 'end'
type Direction = 'top' | 'right' | 'bottom' | 'left'

export class StepNode {
  description: string
  label: string
  value: string = ''
  inputLinks: Link[] = []
  outputLinks: Link[] = []
  id: number
  type: NodeType

  constructor(description: string, label: string, id: number, type: NodeType = 'normal') {
    this.description = description
    this.label = label
    this.id = id
    this.type = type
  }

  static buildStepNodes(stepsInfo: { description: string; label: string }[]): StepNode[] {
    return stepsInfo.map((info, i) => new StepNode(info.description, info.label, i))
  }
}

interface Link {
  source: StepNode
  dest: StepNode
  id: string
}

type ContainerSelection = d3.Selection<d3.ContainerElement, unknown, HTMLElement, any>
type SVGGSelection = d3.Selection<SVGGElement, unknown, HTMLElement, any>
type SVGSVGSlection = d3.Selection<SVGSVGElement, unknown, HTMLElement, any>
type SVGPathSelection = d3.Selection<SVGPathElement, unknown, HTMLElement, any>

const STEP_MARGIN = 30
const STEP_CARD_HEIGHT = 100

const nodeScale = (index: number) => index * (STEP_CARD_HEIGHT + STEP_MARGIN)

export class StepChart {
  private steps: StepNode[]
  private container: ContainerSelection
  private chart: SVGGSelection
  private node2coord: { [id: string]: number } = {}

  constructor(container: ContainerSelection, steps: StepNode[]) {
    this.container = container
    this.steps = steps

    // Remove all svg and div elements
    container.select('svg g.core').remove()
    container.selectAll('div.step-card').remove()

    // Get width and height of svg
    const node = container.node() as HTMLDivElement
    const totalWidth = node.offsetWidth
    const totalHeight = node.offsetHeight
    // const margin = { top: 3, right: 3, bottom: 3, left: 3 }
    // const width = totalWidth - margin.left - margin.right
    // const height = totalHeight - margin.top - margin.bottom

    this.chart = container
      .select('svg')
      .attr('width', totalWidth)
      .attr('height', totalHeight)
      .append('g')
      .classed('core', true)
    // .attr('transform', `translate(${margin.left},${margin.top})`)
  }

  linkNode(sourceId: number, destId: number) {
    const sourceNode = this.steps[sourceId]
    const destNode = this.steps[destId]

    const link: Link = {
      source: sourceNode,
      dest: destNode,
      id: `link-${sourceId}-${destId}`
    }

    sourceNode.outputLinks.push(link)
    destNode.inputLinks.push(link)

    return this
  }

  render() {
    for (const step of this.steps) {
      this.drawStep(step)
    }

    for (const step of this.steps) {
      for (const link of step.outputLinks) {
        this.drawLink(link)
      }
    }

    if (window) {
      window.addEventListener('resize', () => {
        // Redraw links
        this.container.selectAll('.link').remove()
        for (const step of this.steps) {
          for (const link of step.outputLinks) {
            this.drawLink(link)
          }
        }
      })
    }
  }

  update(stepId: number, value: string) {
    // Remove all active class
    this.container
      .selectAll('.v-card--raised')
      .classed('v-card--raised', false)
      .classed('v-card--outlined', true)

    // Add active class to current step card
    const target = this.container
      .select('#step-card-' + stepId)
      .classed('v-card--raised', true)
      .classed('v-card--outlined', false)

    // Update step label
    target.select('.step-card__label').text(this.steps[stepId].label + ': ' + value)
  }

  private drawStep(step: StepNode) {
    const div = this.container.insert('div', ':first-child').attr('id', `step-card-${step.id}`)
    const y = nodeScale(step.id)
    this.node2coord[step.id] = y
    const classList = 'step-card v-card v-sheet theme--light v-card--outlined mx-auto'

    if (step.type === 'start') {
      div
        .classed(
          'step-card v-card v-sheet theme--light step-card--start v-card--raised mx-auto',
          true
        )
        .style('top', '0px')

      div
        .append('div')
        .classed('v-card__title step-card__label', true)
        .text(step.label)
        .style('justify-content', 'center')
    } else if (step.type === 'end') {
      div.classed(classList + ' step-card--end', true).style('bottom', '0px')

      div
        .append('div')
        .classed('v-card__title step-card__label', true)
        .text(step.label)
        .style('justify-content', 'center')
    } else {
      div.classed(classList, true).style('top', y + 'px')

      // show description
      div
        .append('div')
        .classed('v-card__title step-card__label', true)
        .text(step.label + ' ' + step.value)
      div
        .append('div')
        .classed('v-card__text', true)
        .text(step.description)
    }
  }

  private drawLink(link: Link) {
    const sourceNode = link.source
    const destNode = link.dest

    const sourceDirection: Direction = sourceNode.id < destNode.id ? 'bottom' : 'left'
    const destDeirection: Direction = sourceNode.id < destNode.id ? 'top' : 'left'
    const sourcePos = this.getLinkEndPointPosition(sourceNode, sourceDirection)
    const destPos = this.getLinkEndPointPosition(destNode, destDeirection)

    if (sourceNode.id < destNode.id) {
      const data = {
        source: [sourcePos.x, sourcePos.y] as [number, number],
        target: [destPos.x, destPos.y] as [number, number]
      }
      this.chart
        .insert('path', ':first-child')
        .classed('link', true)
        .attr('marker-end', 'url(#markerArrow)')
        .attr('id', `link-${sourceNode.id}-${destNode.id}`)
        .attr('d', _ => d3.linkVertical()(data))
    } else {
      // Add middle point to draw a curve line
      const data = [
        [sourcePos.x, sourcePos.y],
        [sourcePos.x - 60, (sourcePos.y + destPos.y) / 2],
        [destPos.x, destPos.y]
      ] as [number, number][]
      this.chart
        .insert('path', ':first-child')
        .classed('link', true)
        .attr('marker-end', 'url(#markerArrow)')
        .attr('id', `link-${sourceNode.id}-${destNode.id}`)
        .attr('d', _ => d3.line().curve(d3.curveStep)(data))
    }
  }

  private getLinkEndPointPosition(node: StepNode, direction: Direction) {
    let y = this.node2coord[node.id]
    let x = (this.container.node() as HTMLDivElement).offsetWidth / 2
    const divNode = this.container.select('#step-card-' + node.id).node() as HTMLDivElement

    if (direction === 'left') {
      x -= divNode.offsetWidth / 2
      y += divNode.offsetHeight / 2
    } else if (direction === 'bottom') {
      y += divNode.offsetHeight
    } else if (direction === 'right') {
      x += divNode.offsetWidth / 2
      y += divNode.offsetHeight / 2
    }

    return { x, y }
  }
}