// A from-scratch isochromat Bloch simulator for pulseq sequences.
//
// The magnetisation of every spin is evolved through the sequence in the
// rotating frame. The timeline is split at every gradient vertex, every RF
// sample, and every ADC sample (the union is `boundaries`), so within each
// segment the gradients are linear and the RF is ~one raster long. Two regimes:
//
//   * Free precession (no RF in the segment): the effective field is purely
//     longitudinal, so each spin's transverse magnetisation just rotates about
//     z by the gradient phase  dφ = 2π·(g·r)·dt + Δw·dt  and relaxes. Exact for
//     a linear gradient, so these segments can be as long as a whole readout
//     dwell or delay — this is what makes the simulation fast.
//   * Excitation (RF present): the full 3-D effective field
//     Ω = (2π·B1·cosθ, 2π·B1·sinθ, 2π·g·r + Δw) rotates each spin (Rodrigues).
//
// Signal at each ADC sample = Σ_j ρ_j·(Mx+iMy), demodulated by the receiver
// phase (ADC phase + frequency offset). Units: pulseq gradients are Hz/m and RF
// amplitude Hz, positions m, so g·r and B1 are already in Hz.
import type { ParsedSeq } from '../seq/types.ts'
import { reconstruct } from '../seq/reconstruct.ts'
import type { Series } from '../seq/reconstruct.ts'
import type { Phantom } from '../phantom/phantomTypes.ts'

export interface RawSignal {
  /** Real part of the acquired signal, one entry per ADC sample (all readouts concatenated). */
  re: Float64Array
  im: Float64Array
  /** Number of ADC events (readout lines). */
  numReadouts: number
  /** Samples in each readout, length numReadouts. */
  samplesPerReadout: Int32Array
  /** Prefix offsets into re/im, length numReadouts+1. */
  offsets: Int32Array
  maxSamplesPerReadout: number
}

export interface SimProgress {
  fraction: number
  segment: number
  numSegments: number
  samplesDone: number
  numSamples: number
}

export interface SimOptions {
  onProgress?: (p: SimProgress) => void
  /** How often (in segments) to invoke onProgress. */
  progressEvery?: number
}

const TWO_PI = 2 * Math.PI

/** Linear interpolation of a (t,v) polyline at ascending query times (moving cursor, O(n)). */
function sampleAtTimes(series: Series, times: Float64Array): Float64Array {
  const { t, v } = series
  const n = t.length
  const out = new Float64Array(times.length)
  if (n === 0) return out
  let p = 0
  for (let i = 0; i < times.length; i++) {
    const time = times[i]
    if (time <= t[0]) {
      out[i] = v[0]
      continue
    }
    if (time >= t[n - 1]) {
      out[i] = v[n - 1]
      continue
    }
    while (p < n - 1 && t[p + 1] < time) p++
    const ta = t[p]
    const tb = t[p + 1]
    if (tb === ta) {
      out[i] = v[p + 1]
    } else {
      const f = (time - ta) / (tb - ta)
      out[i] = v[p] + f * (v[p + 1] - v[p])
    }
  }
  return out
}

/** One ADC sample's recording target: which readout, which column, and receiver phase. */
interface AdcSampleRec {
  segment: number // record after the segment ending at this sample's time
  readout: number
  col: number
  cosPhi: number // demod: multiply signal by exp(-i φ_rx) = (cosPhi - i sinPhi)... stored as cos/sin of φ_rx
  sinPhi: number
}

interface Schedule {
  boundaries: Float64Array
  midpoints: Float64Array
  dt: Float64Array
  gxMid: Float64Array
  gyMid: Float64Array
  gzMid: Float64Array
  rfMagMid: Float64Array
  rfActive: Uint8Array
  records: AdcSampleRec[]
  numReadouts: number
  samplesPerReadout: Int32Array
  offsets: Int32Array
  numSamples: number
  rfPhase: Series
  rfMag: Series
}

function buildSchedule(seq: ParsedSeq): Schedule {
  const rec = reconstruct(seq)

  // ADC sample schedule, straight from the parsed events (so we keep phase/freq).
  const adcSampleTimes: number[] = []
  const adcSampleReadout: number[] = []
  const adcSampleCol: number[] = []
  const adcSamplePhi: number[] = []
  const samplesPerReadout: number[] = []
  let readout = 0
  let t0 = 0
  for (let bi = 0; bi < seq.blocks.length; bi++) {
    const block = seq.blocks[bi]
    if (block.adcId !== 0) {
      const adc = seq.adcs.get(block.adcId)
      if (adc && adc.num > 0) {
        const start = t0 + adc.delay
        samplesPerReadout.push(adc.num)
        for (let i = 0; i < adc.num; i++) {
          const tk = start + adc.dwell * (i + 0.5)
          // Receiver phase: constant ADC phase offset + frequency offset ramp.
          const phi = adc.phase + TWO_PI * adc.freq * (tk - start)
          adcSampleTimes.push(tk)
          adcSampleReadout.push(readout)
          adcSampleCol.push(i)
          adcSamplePhi.push(phi)
        }
        readout++
      }
    }
    t0 += block.duration
  }
  const numReadouts = readout
  const numSamples = adcSampleTimes.length

  // Union of all event boundaries (gradient vertices, RF samples, ADC samples, endpoints).
  const times = new Set<number>()
  times.add(0)
  times.add(rec.duration)
  for (const s of [rec.gx, rec.gy, rec.gz, rec.rfMag]) {
    for (let i = 0; i < s.t.length; i++) if (!Number.isNaN(s.t[i])) times.add(s.t[i])
  }
  for (const s of rec.rfSpans) {
    times.add(s.start)
    times.add(s.end)
  }
  for (const tk of adcSampleTimes) times.add(tk)
  const boundaries = Float64Array.from(times)
  boundaries.sort()

  const numSegments = boundaries.length - 1
  const midpoints = new Float64Array(numSegments)
  const dt = new Float64Array(numSegments)
  for (let i = 0; i < numSegments; i++) {
    midpoints[i] = 0.5 * (boundaries[i] + boundaries[i + 1])
    dt[i] = boundaries[i + 1] - boundaries[i]
  }

  const gxMid = sampleAtTimes(rec.gx, midpoints)
  const gyMid = sampleAtTimes(rec.gy, midpoints)
  const gzMid = sampleAtTimes(rec.gz, midpoints)
  const rfMagMid = sampleAtTimes(rec.rfMag, midpoints)

  // Mark segments that fall inside an RF pulse span.
  const rfActive = new Uint8Array(numSegments)
  for (const span of rec.rfSpans) {
    for (let i = 0; i < numSegments; i++) {
      if (midpoints[i] > span.start && midpoints[i] < span.end) rfActive[i] = 1
    }
  }

  // Map each ADC sample time to the segment it ends, via a value->index lookup.
  const indexOf = new Map<number, number>()
  for (let i = 0; i < boundaries.length; i++) indexOf.set(boundaries[i], i)
  const records: AdcSampleRec[] = []
  for (let k = 0; k < numSamples; k++) {
    const bIndex = indexOf.get(adcSampleTimes[k])
    if (bIndex === undefined || bIndex === 0) continue
    const phi = adcSamplePhi[k]
    records.push({
      segment: bIndex - 1,
      readout: adcSampleReadout[k],
      col: adcSampleCol[k],
      cosPhi: Math.cos(phi),
      sinPhi: Math.sin(phi),
    })
  }
  // Group records by the segment they fire after.
  records.sort((a, b) => a.segment - b.segment)

  const spr = Int32Array.from(samplesPerReadout)
  const offsets = new Int32Array(numReadouts + 1)
  for (let r = 0; r < numReadouts; r++) offsets[r + 1] = offsets[r] + spr[r]

  return {
    boundaries,
    midpoints,
    dt,
    gxMid,
    gyMid,
    gzMid,
    rfMagMid,
    rfActive,
    records,
    numReadouts,
    samplesPerReadout: spr,
    offsets,
    numSamples,
    rfPhase: rec.rfPhase,
    rfMag: rec.rfMag,
  }
}

export function simulate(seq: ParsedSeq, phantom: Phantom, opts: SimOptions = {}): RawSignal {
  const sched = buildSchedule(seq)
  const { boundaries, midpoints, dt, gxMid, gyMid, gzMid, rfMagMid, rfActive, records } = sched
  const numSegments = boundaries.length - 1
  const ns = phantom.ns
  const { x, y, z, rho, t1, t2, dw } = phantom

  // Magnetisation, initialised at thermal equilibrium (Mz = ρ).
  const Mx = new Float64Array(ns)
  const My = new Float64Array(ns)
  const Mz = new Float64Array(ns)
  for (let j = 0; j < ns; j++) Mz[j] = rho[j]

  const re = new Float64Array(sched.numSamples)
  const im = new Float64Array(sched.numSamples)

  // Fast path when relaxation is uniform across the phantom (our built-ins are).
  let uniformRelax = true
  for (let j = 1; j < ns; j++) {
    if (t1[j] !== t1[0] || t2[j] !== t2[0]) {
      uniformRelax = false
      break
    }
  }
  const r1u = ns > 0 ? 1 / t1[0] : 0
  const r2u = ns > 0 ? 1 / t2[0] : 0

  const progressEvery = opts.progressEvery ?? 256
  let recPtr = 0

  for (let i = 0; i < numSegments; i++) {
    const dti = dt[i]
    if (dti > 0) {
      if (rfActive[i]) {
        // --- Excitation: full 3-D rotation about the effective field ---
        const b1 = rfMagMid[i]
        const phase = sampleSeriesScalar(sched.rfPhase, midpoints[i])
        const ph = Number.isNaN(phase) ? 0 : phase
        const w1 = TWO_PI * b1
        const wx = w1 * Math.cos(ph)
        const wy = w1 * Math.sin(ph)
        const gx = gxMid[i]
        const gy = gyMid[i]
        const gz = gzMid[i]
        const e1 = uniformRelax ? Math.exp(-dti * r1u) : 0
        const e2 = uniformRelax ? Math.exp(-dti * r2u) : 0
        for (let j = 0; j < ns; j++) {
          const wz = TWO_PI * (gx * x[j] + gy * y[j] + gz * z[j]) + dw[j]
          const wmag = Math.sqrt(wx * wx + wy * wy + wz * wz)
          let mx = Mx[j]
          let my = My[j]
          let mz = Mz[j]
          if (wmag > 0) {
            // Rotate by θ = -wmag·dt about n = (wx,wy,wz)/wmag (sign matches free precession).
            const theta = -wmag * dti
            const c = Math.cos(theta)
            const s = Math.sin(theta)
            const inv = 1 / wmag
            const nx = wx * inv
            const ny = wy * inv
            const nz = wz * inv
            const dot = nx * mx + ny * my + nz * mz
            // Rodrigues: m' = m c + (n×m) s + n (n·m)(1-c)
            const crx = ny * mz - nz * my
            const cry = nz * mx - nx * mz
            const crz = nx * my - ny * mx
            const k = dot * (1 - c)
            mx = mx * c + crx * s + nx * k
            my = my * c + cry * s + ny * k
            mz = mz * c + crz * s + nz * k
          }
          const e2j = uniformRelax ? e2 : Math.exp(-dti / t2[j])
          const e1j = uniformRelax ? e1 : Math.exp(-dti / t1[j])
          Mx[j] = mx * e2j
          My[j] = my * e2j
          Mz[j] = mz * e1j + rho[j] * (1 - e1j)
        }
      } else {
        // --- Free precession: rotate about z by the gradient phase, then relax ---
        const ax = TWO_PI * gxMid[i] * dti
        const ay = TWO_PI * gyMid[i] * dti
        const az = TWO_PI * gzMid[i] * dti
        const e1 = uniformRelax ? Math.exp(-dti * r1u) : 0
        const e2 = uniformRelax ? Math.exp(-dti * r2u) : 0
        for (let j = 0; j < ns; j++) {
          const dphi = ax * x[j] + ay * y[j] + az * z[j] + dw[j] * dti
          const c = Math.cos(dphi)
          const s = Math.sin(dphi)
          const mx = Mx[j]
          const my = My[j]
          const e2j = uniformRelax ? e2 : Math.exp(-dti / t2[j])
          const e1j = uniformRelax ? e1 : Math.exp(-dti / t1[j])
          Mx[j] = (mx * c + my * s) * e2j
          My[j] = (-mx * s + my * c) * e2j
          Mz[j] = Mz[j] * e1j + rho[j] * (1 - e1j)
        }
      }
    }

    // Record any ADC samples that fire at the end of this segment.
    while (recPtr < records.length && records[recPtr].segment === i) {
      const r = records[recPtr]
      let sre = 0
      let sim = 0
      for (let j = 0; j < ns; j++) {
        sre += rho[j] * Mx[j]
        sim += rho[j] * My[j]
      }
      // recorded = (sre + i·sim) · exp(-i φ_rx)
      const idx = sched.offsets[r.readout] + r.col
      re[idx] = sre * r.cosPhi + sim * r.sinPhi
      im[idx] = -sre * r.sinPhi + sim * r.cosPhi
      recPtr++
    }

    if (opts.onProgress && (i % progressEvery === 0 || i === numSegments - 1)) {
      opts.onProgress({
        fraction: numSegments > 0 ? (i + 1) / numSegments : 1,
        segment: i + 1,
        numSegments,
        samplesDone: recPtr,
        numSamples: sched.numSamples,
      })
    }
  }

  let maxSamplesPerReadout = 0
  for (let r = 0; r < sched.numReadouts; r++)
    maxSamplesPerReadout = Math.max(maxSamplesPerReadout, sched.samplesPerReadout[r])

  return {
    re,
    im,
    numReadouts: sched.numReadouts,
    samplesPerReadout: sched.samplesPerReadout,
    offsets: sched.offsets,
    maxSamplesPerReadout,
  }
}

/** Point sample of a (t,v) polyline (binary search); used only at RF midpoints. */
function sampleSeriesScalar(series: Series, time: number): number {
  const { t, v } = series
  const n = t.length
  if (n === 0) return NaN
  if (time <= t[0]) return v[0]
  if (time >= t[n - 1]) return v[n - 1]
  let lo = 0
  let hi = n - 1
  while (hi - lo > 1) {
    const mid = (lo + hi) >> 1
    if (t[mid] <= time) lo = mid
    else hi = mid
  }
  const ta = t[lo]
  const tb = t[hi]
  const va = v[lo]
  const vb = v[hi]
  if (Number.isNaN(va) || Number.isNaN(vb)) return Number.isNaN(va) ? vb : va
  if (tb === ta) return vb
  return va + ((time - ta) / (tb - ta)) * (vb - va)
}