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

import * as d3 from 'd3'
import { Point as PointPosition } from './gradient-descent'
import { SingleVarFunction } from './functions'

type NumberScaleLinear = d3.ScaleLinear<number, number>
type ContainerSelection = d3.Selection<d3.ContainerElement, unknown, HTMLElement, any>
type SVGSelection = d3.Selection<SVGGElement, unknown, HTMLElement, any>
type SVGPathSelection = d3.Selection<SVGPathElement, unknown, HTMLElement, any>

const FLOAT_COMPARE_TOLERANCE = 0.000001

interface Point {
  x: number
  y: number
  circle: d3.Selection<SVGCircleElement, unknown, HTMLElement, any>
  lines: d3.Selection<SVGLineElement, unknown, HTMLElement, any>[]
}

interface Context {
  groundFunc: SingleVarFunction
  chart: SVGSelection
  functionLine: SVGPathSelection
  points: Point[]
}

export class FunctionChart {
  private xScale: NumberScaleLinear
  private yScale: NumberScaleLinear
  chart: SVGSelection
  private points: Point[] = []
  private functionLine: SVGPathSelection
  private domain: number[]
  private range: number[]
  private listeners: ((context: Context) => void)[] = []
  private container: ContainerSelection

  constructor(private groundFunc: SingleVarFunction, container: ContainerSelection) {
    this.container = container
    const node = container.node() as HTMLElement

    const totalWidth = node.offsetWidth
    const totalHeight = node.offsetHeight
    const margin = { top: 60, right: 60, bottom: 60, left: 60 }
    const width = totalWidth - margin.left - margin.right
    const height = totalHeight - margin.top - margin.bottom

    this.domain = groundFunc.domain
    this.range = groundFunc.range
    ;[this.xScale, this.yScale] = this.makeScale(width, height)

    // Remove previous chart
    container.select('svg').remove()

    // Render chart
    this.chart = container
      .append('svg')
      .attr('width', totalWidth)
      .attr('height', totalHeight)
      .append('g')
      .attr('transform', `translate(${margin.left}, ${margin.top})`)
    this.functionLine = this.renderChart(width, height)
  }

  performActionOnChart(callback: (context: Context) => void) {
    const context = {
      groundFunc: this.groundFunc,
      chart: this.chart,
      functionLine: this.functionLine,
      points: this.points
    }
    callback(context)

    this.listeners.push(callback)
  }

  // eslint-disable-next-line @typescript-eslint/no-unused-vars
  private renderChart(width: number, height: number) {
    const chart = this.chart
    // Append scale
    chart
      .append('g')
      .call(d3.axisBottom(this.xScale).ticks(10))
      .attr('transform', `translate(0, ${height})`)
    chart.append('g').call(d3.axisLeft(this.yScale).ticks(10))

    // Draw line
    const points: [number, number][] = this.xScale
      .ticks(100)
      .map(x => [this.xScale(x), this.yScale(this.groundFunc.f(x))])

    const path = d3.line().curve(d3.curveBasis)(points)
    return chart
      .append('path')
      .classed('function-line', true)
      .attr('d', _ => path)
      .style('fill', 'none')
      .style('stroke', '#348')
      .style('stroke-width', '1.5px')
  }

  private makeScale(width: number, height: number) {
    const xScale = d3
      .scaleLinear()
      .domain(this.domain)
      .range([0, width])
      .nice()
    const yScale = d3
      .scaleLinear()
      .domain(this.range)
      .range([height, 0])
      .nice()
    return [xScale, yScale]
  }

  getPoints() {
    return this.points
  }

  setPoint(index: number, p: PointPosition) {
    const posX = this.xScale(p.x)
    const posY = this.yScale(p.y)

    // Update related circles
    const target = index === -1 ? this.points[this.points.length - 1] : this.points[index]
    target.circle.attr('transform', `translate(${posX}, ${posY})`)

    // Update related lines
    for (const line of target.lines) {
      const x1 = Number.parseFloat(line.attr('x1'))
      const y1 = Number.parseFloat(line.attr('y1'))

      if (
        Math.abs(x1 - posX) < FLOAT_COMPARE_TOLERANCE &&
        Math.abs(y1 - posY) < FLOAT_COMPARE_TOLERANCE
      ) {
        line.attr('x1', posX)
        line.attr('y1', posY)
      } else {
        line.attr('x2', posX)
        line.attr('y2', posY)
      }
    }
  }

  appendPoint(x: number) {
    const circleId = this.points.length
    const posX = this.xScale(x)
    const posY = this.yScale(this.groundFunc.f(x))
    const circle = this.chart
      .append('circle')
      .attr('id', `circle-${circleId}`)
      .attr('r', 3)
      .attr('transform', `translate(${posX}, ${posY})`)

    this.points.push({ x, y: this.groundFunc.f(x), circle, lines: [] })
  }

  appendPointByRealPosition(posX: number) {
    const circleId = this.points.length
    const x = this.xScale.invert(posX)
    const posY = this.yScale(this.groundFunc.f(x))
    const circle = this.chart
      .append('circle')
      .attr('id', `circle-${circleId}`)
      .attr('r', 3)
      .attr('transform', `translate(${posX}, ${posY})`)

    this.points.push({ x, y: this.groundFunc.f(x), circle, lines: [] })
  }

  drawLineAt(point1Index: number, point2Index: number) {
    this.drawLineBetweenTwoPoints(this.points[point1Index], this.points[point2Index])
  }

  drawTangentAtPoint(pointIndex: number) {
    this.drawTangent(this.points[pointIndex])
  }

  showDiffAt(point1Index: number, point2Index: number) {
    this.showDiffBetweenPoints(this.points[point1Index], this.points[point2Index])
  }

  private showDiffBetweenPoints(point1: Point, point2: Point) {
    // Get cross point location
    const x = Math.min(point1.x, point2.x)
    const y = Math.max(point1.y, point2.y)
    const posX = this.xScale(x)
    const posY = this.yScale(y)

    const diff = this.chart.append('g')
    diff
      .append('circle')
      .attr('id', `diff-point-${x}-${y}`)
      .attr('r', 3)
      .attr('transform', `translate(${posX}, ${posY})`)

    diff
      .append('line')
      .classed('diff-line', true)
      .attr('x1', posX)
      .attr('x2', x === point1.x ? this.xScale(point2.x) : this.xScale(point1.x))
      .attr('y1', posY)
      .attr('y2', posY)

    diff
      .append('line')
      .classed('diff-line', true)
      .attr('y1', posY)
      .attr('y2', y === point1.y ? this.yScale(point2.y) : this.yScale(point1.y))
      .attr('x1', posX)
      .attr('x2', posX)
  }

  private drawLineBetweenTwoPoints(point1: Point, point2: Point) {
    if (point1.x === point2.x && point1.y === point2.y) {
      // Draw tangent
      this.drawTangent(point1)
    } else {
      const slope = (point1.y - point2.y) / (point1.x - point2.x)
      const lineFunc = (x: number) => (x - point1.x) * slope + point1.y
      const line = this.drawStraightLine(lineFunc)

      point1.lines.push(line)
      point2.lines.push(line)
    }
  }

  private drawTangent(point: Point) {
    const slope = this.groundFunc.df(point.x)
    const lineFunc = (x: number) => (x - point.x) * slope + point.y

    point.lines.push(this.drawStraightLine(lineFunc))
  }

  private drawStraightLine(lineFunc: (x: number) => number) {
    const startPoint = [this.xScale(this.domain[0]), this.yScale(lineFunc(this.domain[0]))]
    const endPoint = [this.xScale(this.domain[1]), this.yScale(lineFunc(this.domain[1]))]
    return this.chart
      .append('line')
      .style('stroke', 'black')
      .attr('x1', startPoint[0])
      .attr('y1', startPoint[1])
      .attr('x2', endPoint[0])
      .attr('y2', endPoint[1])
  }
}