// Static figures for Chapters 5 to 8.
import type { ReactNode } from 'react'
import { Arrowheads, Figure, Label, Node, arrow } from './svg.tsx'
import { ink, method, seqColor } from '../plot/palette.ts'

/** The grid, the five-point stencil, and the sparsity it produces. */
export function FivePointStencil() {
  const n = 8
  const s = 20
  const X0 = 42
  const Y0 = 34
  const gx = (i: number) => X0 + i * s
  const gy = (j: number) => Y0 + (n - j) * s
  const ci = 3
  const cj = 5
  const nb: [number, number][] = [
    [ci - 1, cj],
    [ci + 1, cj],
    [ci, cj - 1],
    [ci, cj + 1],
  ]
  return (
    <Figure
      width={716}
      height={230}
      caption="The five-point Laplacian on a uniform grid, N = 8. Boundary nodes (blue) carry known data; each of the (N−1)² interior nodes gives one equation coupling it to its four neighbours."
    >
      {/* grid lines */}
      {Array.from({ length: n + 1 }, (_, i) => (
        <g key={i}>
          <line x1={gx(i)} y1={gy(0)} x2={gx(i)} y2={gy(n)} stroke={ink.grid} strokeWidth={1} />
          <line x1={gx(0)} y1={gy(i)} x2={gx(n)} y2={gy(i)} stroke={ink.grid} strokeWidth={1} />
        </g>
      ))}
      <rect x={gx(0)} y={gy(n)} width={n * s} height={n * s} fill="none" stroke={method.fd2} strokeWidth={2} />

      {/* the stencil arms, drawn under the nodes */}
      {nb.map(([i, j], k) => (
        <line key={k} x1={gx(ci)} y1={gy(cj)} x2={gx(i)} y2={gy(j)} stroke={method.fd2} strokeWidth={2.4} />
      ))}

      {/* nodes */}
      {Array.from({ length: n + 1 }, (_, i) =>
        Array.from({ length: n + 1 }, (_, j) => {
          const bnd = i === 0 || j === 0 || i === n || j === n
          const inStencil = (i === ci && j === cj) || nb.some(([a, b]) => a === i && b === j)
          return (
            <Node
              key={`${i}-${j}`}
              x={gx(i)}
              y={gy(j)}
              r={inStencil ? 4.6 : bnd ? 3.4 : 2.4}
              fill={inStencil ? method.fd2 : bnd ? '#256abf' : '#5a5a53'}
            />
          )
        }),
      )}
      <Label x={gx(ci) - 15} y={gy(cj) - 13} size={11.5} fill={ink.primary} mono>
        −4
      </Label>
      {nb.map(([i, j], k) => (
        <Label
          key={k}
          x={gx(i) + (j === cj ? (i < ci ? -13 : 13) : 0)}
          y={gy(j) + (j === cj ? 0 : j > cj ? -13 : 13)}
          size={11}
          fill={ink.secondary}
          mono
        >
          1
        </Label>
      ))}
      <Label x={gx(n) + 8} y={gy(Math.round(n / 2))} size={11.5} anchor="start" fill={method.fd2}>
        u = g
      </Label>

      {/* the sparsity pattern of A = I⊗T + T⊗I */}
      <g transform="translate(300 20)">
        <Label x={92} y={-4} size={11.5} fill={ink.muted}>
          A = I ⊗ T + T ⊗ I, for N = 6
        </Label>
        {(() => {
          const m = 5
          const K = m * m
          const px = 168 / K
          const cells: ReactNode[] = []
          for (let r = 0; r < K; r++) {
            for (let c = 0; c < K; c++) {
              const d = Math.abs(r - c)
              if (d !== 0 && d !== 1 && d !== m) continue
              if (d === 1 && Math.floor(r / m) !== Math.floor(c / m)) continue
              cells.push(
                <rect
                  key={`${r}-${c}`}
                  x={8 + c * px}
                  y={8 + r * px}
                  width={px}
                  height={px}
                  fill={d === 0 ? method.fd2 : '#1c5cab'}
                  opacity={d === 0 ? 0.95 : 0.65}
                />,
              )
            }
          }
          return cells
        })()}
        <rect x={8} y={8} width={168} height={168} fill="none" stroke={ink.axis} strokeWidth={1} />
        <Label x={92} y={192} size={11} fill={ink.muted}>
          five non-zeros per row, and dense elimination ignores that
        </Label>
      </g>
    </Figure>
  )
}

/** The sine-transform solve, as three steps. */
export function FastPoissonSteps() {
  const boxes = [
    { x: 10, w: 160, t: 'F', s: 'right-hand side on the grid', c: ink.axis },
    { x: 216, w: 176, t: 'F̃ = S F S', s: '2-D sine transform, O(N² log N)', c: method.fast },
    { x: 438, w: 176, t: 'Ũ = F̃ / (λₖ + λₗ)', s: 'a division, one per unknown', c: method.fast },
  ]
  return (
    <Figure
      width={716}
      height={182}
      caption="The fast Poisson solver. The sine vectors diagonalise the five-point matrix exactly, so the solve is a transform, a division and a transform back. The answer is the same one dense elimination gives, to the last bit."
    >
      <Arrowheads id="fps" colors={[ink.muted]} />
      {boxes.map((b, i) => (
        <g key={i}>
          <rect x={b.x} y={30} width={b.w} height={50} rx={6} fill="#201f1e" stroke={b.c} strokeWidth={1.4} />
          <Label x={b.x + b.w / 2} y={49} size={13} fill={ink.primary}>
            {b.t}
          </Label>
          <Label x={b.x + b.w / 2} y={67} size={10.5} fill={ink.muted}>
            {b.s}
          </Label>
          {i < boxes.length - 1 ? (
            <line x1={b.x + b.w + 8} y1={55} x2={boxes[i + 1].x - 8} y2={55} stroke={ink.muted} strokeWidth={1.3} markerEnd={arrow('fps', 0)} />
          ) : null}
        </g>
      ))}
      <rect x={438} y={106} width={176} height={50} rx={6} fill="#201f1e" stroke={method.fast} strokeWidth={1.4} />
      <Label x={526} y={125} size={13} fill={ink.primary}>
        U = S Ũ S
      </Label>
      <Label x={526} y={143} size={10.5} fill={ink.muted}>
        transform back
      </Label>
      <line x1={526} y1={84} x2={526} y2={102} stroke={ink.muted} strokeWidth={1.3} markerEnd={arrow('fps', 0)} />

      <Label x={216} y={131} size={12} anchor="start" fill={ink.secondary}>
        λₖ = 4 sin²(πk / 2N)
      </Label>
      <Label x={216} y={149} size={11} anchor="start" fill={ink.muted}>
        known in closed form
      </Label>
    </Figure>
  )
}

/** Five-point against the compact nine-point. */
export function NinePointStencil() {
  const s = 44
  const draw = (ox: number, oy: number, w: (di: number, dj: number) => string | null, colour: string) => {
    const cells: { x: number; y: number; v: string }[] = []
    for (const dj of [-1, 0, 1]) {
      for (const di of [-1, 0, 1]) {
        const v = w(di, dj)
        if (v !== null) cells.push({ x: ox + di * s, y: oy - dj * s, v })
      }
    }
    return (
      <g>
        {cells.map((c, i) =>
          c.x === ox && c.y === oy ? null : (
            <line key={`l${i}`} x1={ox} y1={oy} x2={c.x} y2={c.y} stroke={colour} strokeWidth={1.6} opacity={0.55} />
          ),
        )}
        {cells.map((c, i) => (
          <g key={`c${i}`}>
            <circle cx={c.x} cy={c.y} r={15} fill="#201f1e" stroke={colour} strokeWidth={1.6} />
            <Label x={c.x} y={c.y} size={12} fill={colour} mono>
              {c.v}
            </Label>
          </g>
        ))}
      </g>
    )
  }
  return (
    <Figure
      width={716}
      height={220}
      caption="The compact nine-point (Mehrstellen) stencil beside the five-point one. Both couple a node only to its immediate neighbours, so both are diagonalised by the same sine vectors and cost the same to solve; on Laplace's equation the nine-point scheme is sixth-order accurate."
    >
      <Label x={168} y={18} size={12.5} weight={600} fill={ink.primary}>
        five-point, order 2
      </Label>
      {draw(168, 116, (di, dj) => (di === 0 && dj === 0 ? '−4' : di === 0 || dj === 0 ? '1' : null), method.fd2)}
      <Label x={168} y={202} size={11.5} fill={ink.muted}>
        ÷ h²
      </Label>

      <line x1={358} y1={40} x2={358} y2={190} stroke={ink.grid} strokeWidth={1} />

      <Label x={548} y={18} size={12.5} weight={600} fill={ink.primary}>
        nine-point, order 6 on Δu = 0
      </Label>
      {draw(548, 116, (di, dj) => (di === 0 && dj === 0 ? '−20' : di === 0 || dj === 0 ? '4' : '1'), method.fd4)}
      <Label x={548} y={202} size={11.5} fill={ink.muted}>
        ÷ 6h²
      </Label>
    </Figure>
  )
}

/** The multigrid V-cycle. */
export function VCycle() {
  const levels = [
    { label: 'h', w: 300 },
    { label: '2h', w: 210 },
    { label: '4h', w: 140 },
    { label: '8h', w: 84 },
  ]
  const y = (k: number) => 34 + k * 42
  const XC = 358
  const xl = (k: number) => XC - levels[k].w / 2
  const xr = (k: number) => XC + levels[k].w / 2
  return (
    <Figure
      width={716}
      height={218}
      caption="One multigrid V-cycle. A few Gauss–Seidel sweeps make the error smooth; a smooth error is represented on a grid twice as coarse at a quarter of the cost, and the argument repeats. The whole cycle costs O(N²) and reduces the error by a factor that does not depend on N."
    >
      <Arrowheads id="vc" colors={[method.iterative, ink.muted]} />
      {levels.map((L, k) => (
        <g key={k}>
          <line x1={xl(k)} y1={y(k)} x2={xr(k)} y2={y(k)} stroke={ink.grid} strokeWidth={1} />
          <Label x={xl(k) - 16} y={y(k)} size={11.5} anchor="end" fill={ink.muted} mono>
            {L.label}
          </Label>
          {(() => {
            const m = Math.max(3, 9 - 2 * k)
            return Array.from({ length: m }, (_, i) => (
              <Node key={i} x={xl(k) + (i * L.w) / (m - 1)} y={y(k)} r={2.6} fill={ink.family} />
            ))
          })()}
        </g>
      ))}

      {/* descending and ascending legs */}
      {[0, 1, 2].map((k) => (
        <line
          key={`d${k}`}
          x1={xl(k) + 26}
          y1={y(k) + 6}
          x2={xl(k + 1) + 20}
          y2={y(k + 1) - 6}
          stroke={method.iterative}
          strokeWidth={2}
          markerEnd={arrow('vc', 0)}
        />
      ))}
      {[2, 1, 0].map((k) => (
        <line
          key={`u${k}`}
          x1={xr(k + 1) - 20}
          y1={y(k + 1) - 6}
          x2={xr(k) - 26}
          y2={y(k) + 6}
          stroke={method.iterative}
          strokeWidth={2}
          markerEnd={arrow('vc', 0)}
        />
      ))}

      <Label x={152} y={98} size={11.5} anchor="end" fill={method.iterative}>
        smooth,
      </Label>
      <Label x={152} y={114} size={11.5} anchor="end" fill={method.iterative}>
        restrict residual
      </Label>
      <Label x={564} y={98} size={11.5} anchor="start" fill={method.iterative}>
        interpolate correction,
      </Label>
      <Label x={564} y={114} size={11.5} anchor="start" fill={method.iterative}>
        smooth again
      </Label>
      <Label x={XC} y={y(3) + 22} size={11.5} fill={ink.secondary}>
        solve exactly: the coarse grid is tiny
      </Label>

      <Label x={XC} y={198} size={11.5} fill={ink.muted}>
        cost of the whole cycle ≈ 4/3 of one fine-grid sweep
      </Label>
    </Figure>
  )
}

/** What smoothing does to the error, as a picture rather than a claim. */
export function SmoothingIdea() {
  const L = 40
  const R = 320
  const mid = 62
  const wig = (amp: number, hf: number, phase: number) => {
    const pts: string[] = []
    for (let i = 0; i <= 140; i++) {
      const t = i / 140
      const v =
        amp * Math.sin(Math.PI * t) * 0.75 +
        hf * (Math.sin(19 * Math.PI * t + phase) * 0.5 + Math.sin(31 * Math.PI * t) * 0.35)
      pts.push(`${(L + t * (R - L)).toFixed(1)} ${(mid - v).toFixed(1)}`)
    }
    return 'M ' + pts.join(' L ')
  }
  return (
    <Figure
      width={716}
      height={140}
      caption="Two Gauss–Seidel sweeps applied to a rough error. What is left is smooth, and a smooth function loses nothing by being carried to a grid twice as coarse. Gauss–Seidel is a poor solver and a good smoother, and multigrid uses only the second property."
    >
      <Arrowheads id="sm" colors={[ink.muted]} />
      <Label x={(L + R) / 2} y={16} size={12} fill={ink.muted}>
        error before
      </Label>
      <line x1={L} y1={mid} x2={R} y2={mid} stroke={ink.grid} strokeWidth={1} />
      <path d={wig(26, 11, 0.4)} fill="none" stroke={method.iterative} strokeWidth={1.7} />

      <line x1={R + 26} y1={mid} x2={R + 66} y2={mid} stroke={ink.muted} strokeWidth={1.4} markerEnd={arrow('sm', 0)} />
      <Label x={R + 46} y={mid - 12} size={11} fill={ink.muted}>
        2 sweeps
      </Label>

      <g transform={`translate(${R + 92} 0)`}>
        <Label x={(L + R) / 2 - L + 8} y={16} size={12} fill={ink.muted}>
          error after
        </Label>
        <line x1={8} y1={mid} x2={8 + (R - L)} y2={mid} stroke={ink.grid} strokeWidth={1} />
        <path
          d={(() => {
            const pts: string[] = []
            for (let i = 0; i <= 140; i++) {
              const t = i / 140
              pts.push(`${(8 + t * (R - L)).toFixed(1)} ${(mid - 26 * Math.sin(Math.PI * t) * 0.75).toFixed(1)}`)
            }
            return 'M ' + pts.join(' L ')
          })()}
          fill="none"
          stroke={method.iterative}
          strokeWidth={1.7}
        />
        {Array.from({ length: 8 }, (_, i) => (
          <Node key={i} x={8 + (i * (R - L)) / 7} y={mid} r={3} fill={seqColor(0.8)} />
        ))}
        <Label x={8 + (R - L) / 2} y={mid + 40} size={11} fill={ink.muted}>
          eight coarse points now suffice
        </Label>
      </g>
    </Figure>
  )
}