// Figures for the iterative-solver chapter: how the residual falls per
// iteration, how the true error falls, and how both compare with a direct
// solve of the same system.
import Plot from '../../plot/Plot.tsx'
import Legend from '../../plot/Legend.tsx'
import { EndLabel, Line, Markers } from '../../plot/Curves.tsx'
import { ink, marker, method, type MethodKey } from '../../plot/palette.ts'
import type { Num } from '../../engine/types.ts'
import { fmtErr } from '../onedim/panels.tsx'

export interface IterOut {
  names: string[]
  N: number
  its: number[]
  /** squared relative residual and squared error: the MATLAB layer hands back
   *  the squares, for the reason its driver explains */
  R2: Num[][]
  E2: Num[][]
  T: Num[][]
  discErr: number
  directMs: number
  unknowns: number
}

const rootOf = (rows: Num[][]): Num[][] =>
  rows.map((r) => r.map((v) => (v == null || !isFinite(v) || v < 0 ? null : Math.sqrt(v))))

// one palette slot per method, in the order the script lists them
const SLOTS: MethodKey[] = ['fd2', 'fd4', 'fast', 'spectral', 'iterative']

function lowHigh(rows: Num[][], floor = 1e-14): [number, number] {
  let lo = Infinity
  let hi = -Infinity
  for (const r of rows) {
    for (const v of r) {
      if (v == null || !isFinite(v) || v <= 0) continue
      if (v < lo) lo = v
      if (v > hi) hi = v
    }
  }
  if (!isFinite(lo)) return [floor, 1]
  return [Math.max(lo / 3, floor), hi * 3]
}

function lastIndex(row: Num[]): number {
  for (let i = row.length - 1; i >= 0; i--) if (row[i] != null && isFinite(row[i] as number)) return i
  return 0
}

export function IterPanels({ data }: { data: IterOut }) {
  const its = data.its
  const R = rootOf(data.R2)
  const E = rootOf(data.E2)
  const xd: [number, number] = [0, its[its.length - 1]]
  const rd = lowHigh(R)
  const ed = lowHigh([...E, [data.discErr]])
  const td = lowHigh(data.T.map((r) => r.map((v) => (v != null && v > 0 ? v : null))))
  td[0] = Math.max(td[0], 0.05)

  const legend = data.names.map((n, k) => ({ label: n, color: method[SLOTS[k]] }))

  return (
    <>
      <div className="plot-title">relative residual against iteration count</div>
      <Plot height={300} xDomain={xd} yDomain={rd} yLog xLabel="iteration" yLabel="‖r‖ / ‖r₀‖">
        {(f) => (
          <>
            {R.map((row, k) => (
              <g key={k}>
                <Line x={its} y={row} f={f} color={method[SLOTS[k]]} />
                {row.length <= 40 && (
                  <Markers x={its} y={row} f={f} color={method[SLOTS[k]]} shape={marker[SLOTS[k]]} r={3} />
                )}
              </g>
            ))}
          </>
        )}
      </Plot>
      <Legend items={legend} />

      <div className="plot-title">
        true error against iteration count; the dashed line is the discretisation error, which is where the exact
        solution of this linear system sits
      </div>
      <Plot
        height={300}
        xDomain={xd}
        yDomain={ed}
        yLog
        xLabel="iteration"
        yLabel="max |error|"
        margin={{ r: 30 }}
        over={(f) => (
          <>
            {E.map((row, k) => {
              const i = lastIndex(row)
              return (
                <EndLabel
                  key={k}
                  f={f}
                  x={its[i]}
                  y={row[i] as number}
                  text={data.names[k]}
                  color={method[SLOTS[k]]}
                  anchor="end"
                  dx={-6}
                  dy={-7}
                />
              )
            })}
          </>
        )}
      >
        {(f) => (
          <>
            <Line x={xd} y={[data.discErr, data.discErr]} f={f} color={ink.muted} dash="6 4" width={1.5} />
            {E.map((row, k) => (
              <Line key={k} x={its} y={row} f={f} color={method[SLOTS[k]]} />
            ))}
          </>
        )}
      </Plot>
      <Legend
        items={[...legend, { label: 'discretisation error', color: ink.muted, dash: '6 4' }]}
      />

      <div className="plot-title">
        true error against milliseconds of solver time; the marked point is the sine-transform direct solve of the same
        system
      </div>
      <Plot height={300} xDomain={td} yDomain={ed} xLog yLog xLabel="milliseconds" yLabel="max |error|">
        {(f) => (
          <>
            {E.map((row, k) => (
              <Line
                key={k}
                x={(data.T[k] ?? []).map((v) => (v == null || v <= 0 ? NaN : v))}
                y={row}
                f={f}
                color={method[SLOTS[k]]}
              />
            ))}
            <Markers
              x={[data.directMs]}
              y={[data.discErr]}
              f={f}
              color={ink.node}
              shape="star"
              r={7}
            />
            <EndLabel
              f={f}
              x={data.directMs}
              y={data.discErr}
              text="direct solve"
              color={ink.secondary}
              dx={-12}
              dy={-10}
              anchor="end"
            />
          </>
        )}
      </Plot>
      <Legend items={[...legend, { label: 'direct solve', color: ink.node, shape: 'cross' }]} />
    </>
  )
}

export function iterSummary(d: IterOut) {
  const mg = d.E2[d.E2.length - 1]
  const i = lastIndex(mg)
  return [
    { label: 'unknowns', value: String(d.unknowns) },
    { label: 'discretisation error', value: fmtErr(d.discErr) },
    { label: 'multigrid cycles used', value: String(d.its[i]) },
  ]
}