// Turn a RawSignal into a displayable RGBA image: the "raw k-space" grid, one
// row per ADC readout (phase-encode line), columns along the readout. For a
// Cartesian sequence this IS the k-space matrix; for anything else it is still
// the raw acquired data laid out by readout. DOM-free so it can be unit-tested.
import type { RawSignal } from '../sim/simulate.ts'

export type KspaceMode = 'log-magnitude' | 'magnitude' | 'phase' | 'real' | 'imaginary'

export interface KspaceImage {
  width: number
  height: number
  rgba: Uint8ClampedArray
  /** Descriptive scale info for a legend. */
  scale: { label: string; min: number; max: number }
}

export const KSPACE_MODES: { id: KspaceMode; label: string }[] = [
  { id: 'log-magnitude', label: 'log magnitude' },
  { id: 'magnitude', label: 'magnitude' },
  { id: 'phase', label: 'phase' },
  { id: 'real', label: 'real' },
  { id: 'imaginary', label: 'imaginary' },
]

function grayscale(t: number, out: Uint8ClampedArray, o: number) {
  const v = Math.round(255 * Math.max(0, Math.min(1, t)))
  out[o] = v
  out[o + 1] = v
  out[o + 2] = v
  out[o + 3] = 255
}

/** Diverging blue-white-red for signed values, t in [-1,1]. */
function diverging(t: number, out: Uint8ClampedArray, o: number) {
  const x = Math.max(-1, Math.min(1, t))
  let r: number, g: number, b: number
  if (x < 0) {
    const f = 1 + x // 0..1
    r = 255 * f
    g = 255 * f
    b = 255
  } else {
    const f = 1 - x
    r = 255
    g = 255 * f
    b = 255 * f
  }
  out[o] = r
  out[o + 1] = g
  out[o + 2] = b
  out[o + 3] = 255
}

/** Cyclic colormap for phase, t in [0,1) → hue wheel. */
function cyclic(t: number, out: Uint8ClampedArray, o: number) {
  const h = (t % 1) * 6
  const c = 1
  const x = 1 - Math.abs((h % 2) - 1)
  let r = 0,
    g = 0,
    b = 0
  if (h < 1) [r, g, b] = [c, x, 0]
  else if (h < 2) [r, g, b] = [x, c, 0]
  else if (h < 3) [r, g, b] = [0, c, x]
  else if (h < 4) [r, g, b] = [0, x, c]
  else if (h < 5) [r, g, b] = [x, 0, c]
  else [r, g, b] = [c, 0, x]
  out[o] = 255 * r
  out[o + 1] = 255 * g
  out[o + 2] = 255 * b
  out[o + 3] = 255
}

export function renderKspace(signal: RawSignal, mode: KspaceMode, dynamicRangeDb = 60): KspaceImage {
  const width = Math.max(1, signal.maxSamplesPerReadout)
  const height = Math.max(1, signal.numReadouts)
  const rgba = new Uint8ClampedArray(width * height * 4)
  const { re, im, offsets, samplesPerReadout } = signal

  // Global scale factors.
  let maxMag = 0
  let maxAbs = 0
  for (let i = 0; i < re.length; i++) {
    const m = Math.hypot(re[i], im[i])
    if (m > maxMag) maxMag = m
    if (Math.abs(re[i]) > maxAbs) maxAbs = Math.abs(re[i])
    if (Math.abs(im[i]) > maxAbs) maxAbs = Math.abs(im[i])
  }
  if (maxMag === 0) maxMag = 1
  if (maxAbs === 0) maxAbs = 1
  const floorDb = -Math.abs(dynamicRangeDb)

  for (let r = 0; r < height; r++) {
    const n = r < samplesPerReadout.length ? samplesPerReadout[r] : 0
    const base = r < offsets.length ? offsets[r] : 0
    for (let c = 0; c < width; c++) {
      const o = (r * width + c) * 4
      if (c >= n) {
        // padding for ragged readouts
        rgba[o] = 12
        rgba[o + 1] = 14
        rgba[o + 2] = 20
        rgba[o + 3] = 255
        continue
      }
      const idx = base + c
      const reV = re[idx]
      const imV = im[idx]
      switch (mode) {
        case 'log-magnitude': {
          const m = Math.hypot(reV, imV)
          const db = 20 * Math.log10(m / maxMag + 1e-12)
          grayscale((db - floorDb) / -floorDb, rgba, o)
          break
        }
        case 'magnitude':
          grayscale(Math.hypot(reV, imV) / maxMag, rgba, o)
          break
        case 'phase':
          cyclic((Math.atan2(imV, reV) + Math.PI) / (2 * Math.PI), rgba, o)
          break
        case 'real':
          diverging(reV / maxAbs, rgba, o)
          break
        case 'imaginary':
          diverging(imV / maxAbs, rgba, o)
          break
      }
    }
  }

  const scale =
    mode === 'phase'
      ? { label: 'rad', min: -Math.PI, max: Math.PI }
      : mode === 'log-magnitude'
        ? { label: 'dB', min: floorDb, max: 0 }
        : mode === 'magnitude'
          ? { label: '', min: 0, max: maxMag }
          : { label: '', min: -maxAbs, max: maxAbs }

  return { width, height, rgba, scale }
}