// Turn a ParsedSeq into plottable timelines, mirroring pulseq's
// Sequence.waveforms_and_times conventions:
//  - trapezoids -> 4 vertices (3 when flat == 0), skipped when empty
//  - arbitrary gradients on the centers raster -> edge points from
//    first/last plus the raster-center samples
//  - extended trapezoids (time shape on raster edges) -> vertices as stored
//  - RF samples at (i+0.5)*rfRaster or the time shape; phase includes the
//    phase/frequency offsets
// Series are piecewise-linear polylines; NaN values break the line.
import type { ParsedSeq } from './types.ts'

export interface Series {
  t: Float64Array
  v: Float64Array
}

export interface BlockSpan {
  index: number
  id: number
  start: number
  duration: number
}

export interface AdcSpan {
  blockIndex: number
  start: number
  end: number
  num: number
  dwell: number
}

export interface RfPulseSpan {
  blockIndex: number
  start: number
  end: number
  /** Absolute time of the pulse center (NaN for v1.4 files) */
  center: number
  use: string
}

export interface Reconstructed {
  duration: number
  blockSpans: BlockSpan[]
  gx: Series
  gy: Series
  gz: Series
  /** |B1| in Hz */
  rfMag: Series
  /** rad, NaN between pulses */
  rfPhase: Series
  adcSpans: AdcSpan[]
  rfSpans: RfPulseSpan[]
  stats: {
    maxGrad: { gx: number; gy: number; gz: number }
    maxSlew: { gx: number; gy: number; gz: number }
    rfCount: number
    adcCount: number
    adcSamples: number
  }
}

const EPS = 1e-9

class SeriesBuilder {
  t: number[] = []
  v: number[] = []
  push(t: number, v: number) {
    this.t.push(t)
    this.v.push(v)
  }
  breakLine() {
    if (this.v.length > 0 && !Number.isNaN(this.v[this.v.length - 1])) {
      this.push(this.t[this.t.length - 1], NaN)
    }
  }
  build(): Series {
    return { t: Float64Array.from(this.t), v: Float64Array.from(this.v) }
  }
}

function wrapToPi(x: number): number {
  const w = ((x + Math.PI) % (2 * Math.PI)) - Math.PI
  return w < -Math.PI ? w + 2 * Math.PI : w
}

export function reconstruct(seq: ParsedSeq): Reconstructed {
  const { rasters } = seq
  const gradBuilders = {
    gx: new SeriesBuilder(),
    gy: new SeriesBuilder(),
    gz: new SeriesBuilder(),
  }
  const rfMag = new SeriesBuilder()
  const rfPhase = new SeriesBuilder()
  const blockSpans: BlockSpan[] = []
  const adcSpans: AdcSpan[] = []
  const rfSpans: RfPulseSpan[] = []

  const channels = ['gx', 'gy', 'gz'] as const

  // Zero anchors so empty channels still draw a baseline
  for (const ch of channels) gradBuilders[ch].push(0, 0)
  rfMag.push(0, 0)

  let t0 = 0
  for (let bi = 0; bi < seq.blocks.length; bi++) {
    const block = seq.blocks[bi]
    blockSpans.push({ index: bi, id: block.id, start: t0, duration: block.duration })

    for (const ch of channels) {
      const gid = ch === 'gx' ? block.gxId : ch === 'gy' ? block.gyId : block.gzId
      if (gid === 0) continue
      const grad = seq.grads.get(gid)
      if (!grad) continue
      const b = gradBuilders[ch]
      const gt0 = t0 + grad.delay
      if (grad.kind === 'trap') {
        if (Math.abs(grad.flat) > EPS) {
          b.push(gt0, 0)
          b.push(gt0 + grad.rise, grad.amp)
          b.push(gt0 + grad.rise + grad.flat, grad.amp)
          b.push(gt0 + grad.rise + grad.flat + grad.fall, 0)
        } else if (Math.abs(grad.rise) > EPS && Math.abs(grad.fall) > EPS) {
          b.push(gt0, 0)
          b.push(gt0 + grad.rise, grad.amp)
          b.push(gt0 + grad.rise + grad.fall, 0)
        }
        // else: empty gradient, skip
      } else {
        const shape = seq.shapes.get(grad.shapeId)
        if (!shape) continue
        const n = shape.numSamples
        const raster = rasters.gradient
        let tt: Float64Array
        if (grad.timeShapeId === 0) {
          tt = new Float64Array(n)
          for (let i = 0; i < n; i++) tt[i] = (i + 0.5) * raster
        } else {
          const timeShape = seq.shapes.get(grad.timeShapeId)
          if (!timeShape) continue
          tt = new Float64Array(n)
          for (let i = 0; i < n; i++) tt[i] = timeShape.samples[i] * raster
        }
        const onCenters = Math.abs(tt[0] - 0.5 * raster) < 1e-6 * raster
        if (onCenters) {
          // v1.4 files carry no first/last; fall back to the edge samples
          const first = Number.isNaN(grad.first) ? grad.amp * shape.samples[0] : grad.first
          const last =
            Number.isNaN(grad.last) ? grad.amp * shape.samples[n - 1] : grad.last
          const shapeDur = Math.ceil((tt[n - 1] - EPS) / raster) * raster
          b.push(gt0, first)
          for (let i = 0; i < n; i++) b.push(gt0 + tt[i], grad.amp * shape.samples[i])
          b.push(gt0 + shapeDur, last)
        } else {
          // extended trapezoid: vertices as stored
          for (let i = 0; i < n; i++) b.push(gt0 + tt[i], grad.amp * shape.samples[i])
        }
      }
    }

    if (block.rfId !== 0) {
      const ev = seq.rf.get(block.rfId)
      const magShape = ev ? seq.shapes.get(ev.magShapeId) : undefined
      if (ev && magShape) {
        const n = magShape.numSamples
        const raster = rasters.rf
        const phaseShape = ev.phaseShapeId !== 0 ? seq.shapes.get(ev.phaseShapeId) : undefined
        let tt: Float64Array
        if (ev.timeShapeId === 0) {
          tt = new Float64Array(n)
          for (let i = 0; i < n; i++) tt[i] = (i + 0.5) * raster
        } else {
          const timeShape = seq.shapes.get(ev.timeShapeId)
          tt = new Float64Array(n)
          for (let i = 0; i < n; i++) tt[i] = (timeShape?.samples[i] ?? 0) * raster
        }
        const rt0 = t0 + ev.delay
        const end = rt0 + Math.ceil((tt[n - 1] - EPS) / raster) * raster
        rfMag.push(rt0, 0)
        rfPhase.breakLine()
        for (let i = 0; i < n; i++) {
          const t = rt0 + tt[i]
          rfMag.push(t, Math.abs(ev.amp * magShape.samples[i]))
          const ph = 2 * Math.PI * (phaseShape?.samples[i] ?? 0) + ev.phase + 2 * Math.PI * ev.freq * tt[i]
          rfPhase.push(t, wrapToPi(ph))
        }
        rfPhase.breakLine()
        rfMag.push(end, 0)
        rfSpans.push({
          blockIndex: bi,
          start: rt0,
          end,
          center: Number.isNaN(ev.center) ? NaN : rt0 + ev.center,
          use: ev.use,
        })
      }
    }

    if (block.adcId !== 0) {
      const adc = seq.adcs.get(block.adcId)
      if (adc) {
        const start = t0 + adc.delay
        adcSpans.push({
          blockIndex: bi,
          start,
          end: start + adc.num * adc.dwell,
          num: adc.num,
          dwell: adc.dwell,
        })
      }
    }

    t0 += block.duration
  }

  for (const ch of channels) gradBuilders[ch].push(t0, 0)
  rfMag.push(t0, 0)

  const stats = {
    maxGrad: { gx: 0, gy: 0, gz: 0 },
    maxSlew: { gx: 0, gy: 0, gz: 0 },
    rfCount: rfSpans.length,
    adcCount: adcSpans.length,
    adcSamples: adcSpans.reduce((acc, a) => acc + a.num, 0),
  }
  const series = {
    gx: gradBuilders.gx.build(),
    gy: gradBuilders.gy.build(),
    gz: gradBuilders.gz.build(),
  }
  for (const ch of channels) {
    const { t, v } = series[ch]
    for (let i = 0; i < v.length; i++) {
      if (Number.isNaN(v[i])) continue
      stats.maxGrad[ch] = Math.max(stats.maxGrad[ch], Math.abs(v[i]))
      if (i > 0 && !Number.isNaN(v[i - 1]) && t[i] > t[i - 1] + EPS) {
        stats.maxSlew[ch] = Math.max(stats.maxSlew[ch], Math.abs((v[i] - v[i - 1]) / (t[i] - t[i - 1])))
      }
    }
  }

  return {
    duration: t0,
    blockSpans,
    ...series,
    rfMag: rfMag.build(),
    rfPhase: rfPhase.build(),
    adcSpans,
    rfSpans,
    stats,
  }
}