/// ===========================================================================
/// THE ENTROPY LAYER — everything between the ONNX tensors and the rANS coder.
/// ===========================================================================
///
/// AEIC's four-stage masked context model, ported from
/// `aic/aeic/src/codec/codec_practical.py` (`PixelCodec.compress` /
/// `.decompress`). This file owns the parts that are *not* neural and *not*
/// range coding:
///
///   * the four checkerboard masks (`get_mask_four_parts`),
///   * `sequeeze` / `unsequeeze_with_mask` (the 4:1 channel fold),
///   * `torch.round`'s half-to-**even** tie rule,
///   * `my_build_indexes` (scale -> CDF row selector),
///   * the z index array (`_build_indexes`: the channel arange),
///   * the fixed call order z, y0, y1, y2, y3.
///
/// It is deliberately pure Dart: no ONNX, no dart:io, no Flutter. The neural
/// half arrives through [AeicEntropyNetwork] and the range coder through
/// [AeicRansCoderFactory], both injected. That is what makes the arithmetic
/// unit-testable without a 67 MB graph or a model download.
///
/// ## Why every detail here is load-bearing
///
/// The rANS decoder re-derives symbol probabilities by re-running the same
/// network and the same arithmetic. A single wrong symbol *position* — an
/// off-by-one in a mask, the wrong channel-group permutation, `sequeeze`
/// folding in the wrong order — does not raise: it desynchronises the coder and
/// produces a sharp, plausible, wrong image. Same for a rounding tie resolved
/// away from zero instead of to even. Both failure modes are silent.
///
/// The golden vectors in `test/services/image_codec_entropy_test.dart` were
/// generated by `aic/exp/export_entropy_layer_golden.py`, which runs the real
/// torch ops on a closed-form input both languages reproduce bit-for-bit.
library;

import 'dart:math' as math;
import 'dart:typed_data';

/// CDF group indices, fixed by `EntropyCoder.update()`'s `add_cdf` call order:
/// the entropy bottleneck (z) is registered first, the Gaussian conditional
/// (y) second. The group index is part of the bitstream format.
const int kAeicZCdfGroup = 0;
const int kAeicYCdfGroup = 1;

/// `my_build_indexes` constants (`codec_practical.py`, `SCALES_MIN/MAX/LEVELS`).
const int kAeicScalesLevels = 64;
const double kAeicLogScaleMin = -2.2072749131897207; // ln(0.11)
const double kAeicLogScaleStep = 0.12305479932808384; // (ln(256)-ln(0.11))/63
const double kAeicScaleThreshold = 0.08; // below this the symbol is skipped
const double kAeicScaleFloor = 1e-5; // torch.maximum floor applied first

/// Tensor shapes for one image, derived the way `compress()` derives them.
///
/// `y` is the image downsampled by 32; `z` is `y` downsampled by a further 4,
/// rounded **up** (`compress()` reflect-pads `y` to a multiple of 4 first).
class AeicEntropyGeometry {
  final int resolution;

  /// Latent channel count `M` (256 for AEIC-SE, which is what ships).
  final int yChannels;
  final int yHeight;
  final int yWidth;

  /// Hyper-latent channels: `M // 2`.
  final int zChannels;
  final int zHeight;
  final int zWidth;

  const AeicEntropyGeometry._({
    required this.resolution,
    required this.yChannels,
    required this.yHeight,
    required this.yWidth,
    required this.zChannels,
    required this.zHeight,
    required this.zWidth,
  });

  factory AeicEntropyGeometry.forResolution(
    int resolution, {
    int yChannels = 256,
  }) {
    if (resolution <= 0 || resolution % 32 != 0) {
      throw ArgumentError.value(
        resolution,
        'resolution',
        'must be a positive multiple of 32 (g_a downsamples by 32)',
      );
    }
    if (yChannels % 4 != 0) {
      throw ArgumentError.value(
        yChannels,
        'yChannels',
        'the four-part mask splits the channels into 4 equal groups',
      );
    }
    final yh = resolution ~/ 32;
    final zh = (yh + 3) ~/ 4;
    return AeicEntropyGeometry._(
      resolution: resolution,
      yChannels: yChannels,
      yHeight: yh,
      yWidth: yh,
      zChannels: yChannels ~/ 2,
      zHeight: zh,
      zWidth: zh,
    );
  }

  /// Channels after `sequeeze` folds the four groups together.
  int get squeezedChannels => yChannels ~/ 4;

  int get yElements => yChannels * yHeight * yWidth;

  /// Symbols coded per y stage (`[1, M/4, y_h, y_w]` flattened).
  int get symbolsPerStage => squeezedChannels * yHeight * yWidth;

  int get zElements => zChannels * zHeight * zWidth;

  /// Total entries handed to the coder: z once, then four y stages.
  int get totalEntries => zElements + 4 * symbolsPerStage;

  List<int> get yShape => <int>[1, yChannels, yHeight, yWidth];

  List<int> get zShape => <int>[1, zChannels, zHeight, zWidth];

  List<int> get imageShape => <int>[1, 3, resolution, resolution];
}

/// The four masks of `get_mask_four_parts`, as arithmetic rather than tensors.
///
/// Reading the Python: each mask is four channel-groups of `M/4` channels
/// stacked, and each group carries one of the four 2x2 micro-patterns
///
///     micro 0 = (y%2, x%2) == (0, 0)      micro 1 = (0, 1)
///     micro 2 = (1, 0)                    micro 3 = (1, 1)
///
/// so `micro k` is live exactly where `(y%2)*2 + (x%2) == k`. The stacking order
/// per stage is
///
///     mask_0 = [m0, m1, m2, m3]     mask_1 = [m3, m2, m1, m0]
///     mask_2 = [m2, m3, m0, m1]     mask_3 = [m1, m0, m3, m2]
///
/// which is exactly `micro = group XOR perm[stage]` with `perm = [0, 3, 2, 1]`.
/// Verified against the tensors themselves in the golden test, because getting
/// this permutation wrong is the single easiest way to silently reorder every
/// symbol in the stream.
class AeicMaskSet {
  /// `perm[stage]`: XOR it with the channel group to get the micro-pattern.
  static const List<int> stagePermutation = <int>[0, 3, 2, 1];

  final AeicEntropyGeometry geometry;

  const AeicMaskSet(this.geometry);

  /// Micro-pattern index used by [channelGroup] in [stage].
  static int microFor(int stage, int channelGroup) =>
      channelGroup ^ stagePermutation[stage];

  /// The single channel group that is live at `(y, x)` in [stage].
  ///
  /// Inverse of [microFor]: the position selects the micro-pattern, and exactly
  /// one group carries it.
  static int liveGroupAt(int stage, int y, int x) =>
      (((y & 1) << 1) | (x & 1)) ^ stagePermutation[stage];

  bool isLive(int stage, int channel, int y, int x) {
    final group = channel ~/ geometry.squeezedChannels;
    return microFor(stage, group) == (((y & 1) << 1) | (x & 1));
  }

  /// Materialises `mask_i` as `[1, M, y_h, y_w]` float32 — only needed when a
  /// caller wants to hand the mask to something else (or check it in a test).
  Float32List maskTensor(int stage) {
    final out = Float32List(geometry.yElements);
    final h = geometry.yHeight;
    final w = geometry.yWidth;
    for (var c = 0; c < geometry.yChannels; c++) {
      final base = c * h * w;
      for (var y = 0; y < h; y++) {
        for (var x = 0; x < w; x++) {
          if (isLive(stage, c, y, x)) {
            out[base + y * w + x] = 1.0;
          }
        }
      }
    }
    return out;
  }

  /// `sequeeze`: sum the four channel chunks, `(g0 + g1) + (g2 + g3)`.
  ///
  /// The association matters in principle (float addition is not associative)
  /// so it is reproduced exactly, even though in practice three of the four
  /// terms are a hard zero for a masked tensor.
  Float32List squeeze(Float32List full) {
    if (full.length != geometry.yElements) {
      throw ArgumentError.value(
        full.length,
        'full',
        'expected ${geometry.yElements} elements (${geometry.yShape})',
      );
    }
    final plane = geometry.yHeight * geometry.yWidth;
    final cq = geometry.squeezedChannels;
    final out = Float32List(cq * plane);
    for (var c = 0; c < cq; c++) {
      final o = c * plane;
      final g0 = c * plane;
      final g1 = (c + cq) * plane;
      final g2 = (c + 2 * cq) * plane;
      final g3 = (c + 3 * cq) * plane;
      for (var i = 0; i < plane; i++) {
        // Each partial sum is forced back through float32, matching torch's
        // `(a + b) + (c + d)` on float32 tensors.
        final left = f32(full[g0 + i] + full[g1 + i]);
        final right = f32(full[g2 + i] + full[g3 + i]);
        out[o + i] = left + right;
      }
    }
    return out;
  }

  /// `unsequeeze_with_mask`: broadcast a squeezed tensor back to `[1, M, h, w]`,
  /// keeping each element only in the channel group whose mask is live there.
  Float32List unsqueeze(Float32List squeezed, int stage) {
    final plane = geometry.yHeight * geometry.yWidth;
    final cq = geometry.squeezedChannels;
    if (squeezed.length != cq * plane) {
      throw ArgumentError.value(
        squeezed.length,
        'squeezed',
        'expected ${cq * plane} elements',
      );
    }
    final out = Float32List(geometry.yElements);
    final w = geometry.yWidth;
    for (var c = 0; c < cq; c++) {
      for (var y = 0; y < geometry.yHeight; y++) {
        for (var x = 0; x < w; x++) {
          final group = liveGroupAt(stage, y, x);
          out[(group * cq + c) * plane + y * w + x] =
              squeezed[c * plane + y * w + x];
        }
      }
    }
    return out;
  }

  /// `t * mask_i`, elementwise, into a fresh tensor.
  Float32List applyMask(Float32List full, int stage) {
    final out = Float32List(full.length);
    final plane = geometry.yHeight * geometry.yWidth;
    final cq = geometry.squeezedChannels;
    final w = geometry.yWidth;
    for (var c = 0; c < geometry.yChannels; c++) {
      final group = c ~/ cq;
      for (var y = 0; y < geometry.yHeight; y++) {
        for (var x = 0; x < w; x++) {
          if (microFor(stage, group) == (((y & 1) << 1) | (x & 1))) {
            final i = c * plane + y * w + x;
            out[i] = full[i];
          }
        }
      }
    }
    return out;
  }

  /// `base * (1 - mask_i) + stageLatent`, the context update between stages.
  ///
  /// `stageLatent` must already be masked (it is the output of [unsqueeze]), so
  /// this is a select, not an add — which is what keeps it exact in float32.
  Float32List mergeContext(
    Float32List base,
    Float32List stageLatent,
    int stage,
  ) {
    final out = Float32List.fromList(base);
    final plane = geometry.yHeight * geometry.yWidth;
    final cq = geometry.squeezedChannels;
    final w = geometry.yWidth;
    for (var c = 0; c < geometry.yChannels; c++) {
      final group = c ~/ cq;
      for (var y = 0; y < geometry.yHeight; y++) {
        for (var x = 0; x < w; x++) {
          if (microFor(stage, group) == (((y & 1) << 1) | (x & 1))) {
            final i = c * plane + y * w + x;
            out[i] = stageLatent[i];
          }
        }
      }
    }
    return out;
  }
}

/// Scratch cell used to force a float64 value through float32 rounding.
final Float32List _f32Cell = Float32List(1);

/// Rounds [value] to float32 exactly as a store into a float32 tensor would.
double f32(double value) {
  _f32Cell[0] = value;
  return _f32Cell[0];
}

/// `torch.round`: ties go to the **even** neighbour, not away from zero.
///
/// Dart's `double.round()` and `roundToDouble()` both round half away from
/// zero, so `(-0.5).round() == -1` where torch gives `0`. On the ft32 corpus
/// `y - means` lands on a tie often enough that this is not theoretical.
double roundHalfToEven(double value) {
  if (!value.isFinite) {
    return value;
  }
  final floor = value.floorToDouble();
  final diff = value - floor;
  if (diff > 0.5) {
    return floor + 1.0;
  }
  if (diff < 0.5) {
    return floor;
  }
  // Exactly halfway: pick the even neighbour.
  return floor % 2.0 == 0.0 ? floor : floor + 1.0;
}

/// `my_build_indexes` — scale -> CDF row selector, in float32, per stage.
///
/// ```
/// s   = max(scale, 1e-5)
/// q   = (ln(s) - log_scale_min) / log_scale_step        # all float32
/// idx = (s < 0.08) ? -1 : trunc(clamp(q, 0, 63))        # trunc, not round
/// ```
///
/// Every intermediate is forced back to float32 because torch evaluates this on
/// a float32 tensor with float32 scalars, and the boundary margin measured in
/// `aic/results/bitexact_encoder.md` is only 1.12x — the tightest number in the
/// whole system. Dart's `log` is still not guaranteed bit-identical to ORT's
/// float32 `Log`; see the note in `aic/results/rans_port_spec.md` §8, which
/// recommends the graph emit these indexes itself. [aeicBuildIndexes] is the
/// fallback, and both sides of the channel run this same code, so sender and
/// receiver agree with each other regardless.
Int16List aeicBuildIndexes(Float32List scales) {
  final logMin = f32(kAeicLogScaleMin);
  final logStep = f32(kAeicLogScaleStep);
  final floor = f32(kAeicScaleFloor);
  final threshold = f32(kAeicScaleThreshold);
  final out = Int16List(scales.length);
  for (var i = 0; i < scales.length; i++) {
    final raw = scales[i];
    final s = raw > floor ? raw : floor;
    if (s < threshold) {
      out[i] = -1;
      continue;
    }
    var q = f32(f32(f32(math.log(s)) - logMin) / logStep);
    if (q < 0.0) {
      q = 0.0;
    } else if (q > kAeicScalesLevels - 1) {
      q = (kAeicScalesLevels - 1).toDouble();
    }
    out[i] = q.toInt(); // truncation toward zero, matching Tensor.int()
  }
  return out;
}

/// The z index array: `_build_indexes` broadcasts `arange(C)` over `H x W`, so
/// this is `H*W` copies of each channel index. Always >= 0.
Int16List aeicZIndexes(AeicEntropyGeometry geometry) {
  final plane = geometry.zHeight * geometry.zWidth;
  final out = Int16List(geometry.zElements);
  var i = 0;
  for (var c = 0; c < geometry.zChannels; c++) {
    for (var p = 0; p < plane; p++) {
      out[i++] = c;
    }
  }
  return out;
}

/// float32 tensor of integers -> int16 symbols, asserting the cast is lossless.
///
/// `py_rans` force-casts through `py::array_t<int16_t>`; the golden vectors
/// store the post-cast values. Anything that does not survive the cast means
/// the entropy model has gone somewhere the format cannot represent, and
/// truncating quietly would desync the decoder.
Int16List aeicToSymbols(Float32List values) {
  final out = Int16List(values.length);
  for (var i = 0; i < values.length; i++) {
    final v = values[i];
    final n = v.toInt();
    if (n != v || n < -32768 || n > 32767) {
      throw StateError(
        'entropy symbol $v at index $i is not a lossless int16; the entropy '
        'model produced a value the bitstream format cannot carry',
      );
    }
    out[i] = n;
  }
  return out;
}

/// RGB bytes -> the graph's `image` input, reproducing torchvision's
/// `ToTensor()` + `Normalize([0.5]*3, [0.5]*3)`: `x = (b/255 - 0.5) / 0.5`.
Float32List aeicRgbToChw(Uint8List rgb, int resolution) {
  final pixels = resolution * resolution;
  if (rgb.length != pixels * 3) {
    throw ArgumentError.value(
      rgb.length,
      'rgb',
      'expected ${pixels * 3} bytes for ${resolution}x$resolution RGB',
    );
  }
  final out = Float32List(pixels * 3);
  for (var c = 0; c < 3; c++) {
    final plane = c * pixels;
    for (var i = 0; i < pixels; i++) {
      final v = f32(rgb[i * 3 + c] / 255.0);
      out[plane + i] = f32(v - 0.5) * 2.0; // /0.5 is exact
    }
  }
  return out;
}

/// Everything the send side needs out of one forward pass of the entropy graph.
///
/// Matches the outputs of `aic/onnx/aeic_entropy_side_fp32_op17.onnx`, which
/// runs `g_a`, `h_a`, `h_s`, `g_c` and all four adapters in one shot — the
/// encoder never has to be incremental, only the decoder does.
class AeicEncodeSideTensors {
  /// `round(z - z_offset)`, `[1, M/2, z_h, z_w]`, integral floats.
  final Float32List zQ;

  /// `round(y * mask_i - means_i)` per stage, `[1, M, y_h, y_w]`, integral.
  final List<Float32List> yQ;

  /// `scales_supp_i * mask_i` per stage, `[1, M, y_h, y_w]`.
  final List<Float32List> scales;

  const AeicEncodeSideTensors({
    required this.zQ,
    required this.yQ,
    required this.scales,
  });
}

/// One decode stage's entropy parameters, **before** masking.
///
/// `adapter_out[i](g_c(adapter_in[i](base)))` split in half; the caller applies
/// `mask_i`. Kept unmasked so the graph does not have to know the stage index
/// at export time.
class AeicStageParams {
  final Float32List meansSupp;
  final Float32List scalesSupp;

  const AeicStageParams({required this.meansSupp, required this.scalesSupp});
}

/// The neural half of the entropy path, as this file needs it.
///
/// ## Graph contract
///
/// **Send side** (exists today, `aeic_entropy_side_fp32_op17.onnx`, 67 MB):
/// input `image [1,3,512,512]` -> outputs `z_q`, `yq0..yq3`, `sc0..sc3`.
///
/// **Receive side** (NOT exported yet — see the report for task B2): the
/// decoder cannot use the send-side graph, because it must interleave network
/// evaluation with symbol decoding: stage `i`'s indexes are unknown until
/// stages `< i` have been decoded and fed back through `g_c`. It needs
///
///   * `z_q [1,128,4,4]` -> `base0 [1,256,16,16]` (`h_s` with `z_offset` baked
///     in, exactly as the send-side export bakes it), and
///   * `base [1,256,16,16]` -> `means{i}`, `scales{i}` for `i` in 0..3
///     (unmasked `adapter_out[i](g_c(adapter_in[i](base)))`, split in half).
///
/// Both can live in one graph with two inputs and nine outputs; ORT prunes to
/// the requested output set, so asking only for `base0` does not run the
/// adapters and asking only for `means2`/`scales2` does not run `h_s`.
abstract class AeicEntropyNetwork {
  /// Whether [runHyperSynthesis] and [runStage] are available. False for a
  /// graph that only carries the send-side path.
  bool get supportsDecodeSide;

  Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw);

  /// `h_s(z_q + z_offset)[:, :, :y_h, :y_w]`.
  Future<Float32List> runHyperSynthesis(Float32List zQ);

  /// `adapter_out[stage](g_c(adapter_in[stage](base)))`, unmasked.
  Future<AeicStageParams> runStage(int stage, Float32List base);
}

/// The rANS encoder, as the entropy layer uses it.
///
/// Implemented by the pure-Dart coder (task B1). Deliberately narrow: five
/// [pushSymbols] calls in the fixed order z, y0, y1, y2, y3, then one [finish].
abstract class AeicRansEncoder {
  /// Appends one `(symbols, indexes)` pair. `indexes[i] < 0` emits nothing.
  void pushSymbols(Int16List symbols, Int16List indexes, int cdfGroup);

  /// Flushes both sub-streams and returns the container bytes.
  Uint8List finish();
}

/// The rANS decoder, as the entropy layer uses it.
///
/// **Must be incremental**: each call resumes the rANS state where the last one
/// stopped. An implementation that decodes the whole stream up front cannot
/// work here — see the class docs on [AeicEntropyNetwork].
abstract class AeicRansDecoder {
  Int16List decodeStream(Int16List indexes, int cdfGroup);
}

/// Supplies coders bound to a parsed CDF table set (task B1 owns both).
abstract class AeicRansCoderFactory {
  AeicRansEncoder createEncoder();

  AeicRansDecoder createDecoder(Uint8List bitstream);
}

/// Thrown when the entropy loop is asked to run without a piece it needs.
class AeicEntropyUnavailable implements Exception {
  final String detail;

  const AeicEntropyUnavailable(this.detail);

  @override
  String toString() => 'AEIC entropy path unavailable: $detail';
}

/// Raised by [ImageCodecBackend] callers that pass a cancel predicate.
typedef AeicCancelCheck = bool Function();

/// The four-stage masked encode and decode, orchestrating [AeicEntropyNetwork]
/// and [AeicRansCoderFactory].
class AeicEntropyCodec {
  final AeicEntropyGeometry geometry;
  final AeicMaskSet masks;
  final AeicEntropyNetwork network;
  final AeicRansCoderFactory coders;

  AeicEntropyCodec({
    required this.geometry,
    required this.network,
    required this.coders,
  }) : masks = AeicMaskSet(geometry);

  /// Packed 8-bit RGB -> rANS bitstream (payload only, no chunk headers).
  ///
  /// One forward pass, then z, y0, y1, y2, y3 into the coder in that order.
  Future<Uint8List> encode(
    Uint8List rgbBytes, {
    void Function(double progress)? onProgress,
    AeicCancelCheck? shouldCancel,
  }) async {
    onProgress?.call(0.02);
    final input = aeicRgbToChw(rgbBytes, geometry.resolution);
    _checkCancel(shouldCancel);

    final tensors = await network.runEncodeSide(input);
    onProgress?.call(0.80);
    _checkCancel(shouldCancel);

    if (tensors.zQ.length != geometry.zElements) {
      throw StateError(
        'entropy graph returned ${tensors.zQ.length} z values, expected '
        '${geometry.zElements} (${geometry.zShape})',
      );
    }
    if (tensors.yQ.length != 4 || tensors.scales.length != 4) {
      throw StateError(
        'entropy graph must return four y_q and four scales tensors, got '
        '${tensors.yQ.length} / ${tensors.scales.length}',
      );
    }

    final encoder = coders.createEncoder();
    encoder.pushSymbols(
      aeicToSymbols(tensors.zQ),
      aeicZIndexes(geometry),
      kAeicZCdfGroup,
    );
    for (var stage = 0; stage < 4; stage++) {
      final symbols = aeicToSymbols(masks.squeeze(tensors.yQ[stage]));
      final indexes = aeicBuildIndexes(masks.squeeze(tensors.scales[stage]));
      encoder.pushSymbols(symbols, indexes, kAeicYCdfGroup);
      onProgress?.call(0.80 + 0.04 * (stage + 1));
      _checkCancel(shouldCancel);
    }
    final stream = encoder.finish();
    onProgress?.call(1.0);
    return stream;
  }

  /// rANS bitstream -> `y_hat [1, M, y_h, y_w]`, ready for the synthesis graph.
  ///
  /// The mirror of [encode], and necessarily incremental: each stage's indexes
  /// come from scales that only exist once the previous stage's symbols have
  /// been decoded and pushed back through the context model.
  Future<Float32List> decodeToLatent(
    Uint8List bitstream, {
    void Function(double progress)? onProgress,
    AeicCancelCheck? shouldCancel,
  }) async {
    if (!network.supportsDecodeSide) {
      throw const AeicEntropyUnavailable(
        'the installed entropy graph is send-side only: it maps image -> '
        'symbols and has no z_q -> base0 / base -> means,scales entry points, '
        'which decoding requires',
      );
    }
    onProgress?.call(0.02);
    final decoder = coders.createDecoder(bitstream);

    final zSymbols = decoder.decodeStream(
      aeicZIndexes(geometry),
      kAeicZCdfGroup,
    );
    final zQ = Float32List(zSymbols.length);
    for (var i = 0; i < zSymbols.length; i++) {
      zQ[i] = zSymbols[i].toDouble();
    }
    _checkCancel(shouldCancel);

    var base = await network.runHyperSynthesis(zQ);
    if (base.length != geometry.yElements) {
      throw StateError(
        'hyper synthesis returned ${base.length} values, expected '
        '${geometry.yElements} (${geometry.yShape})',
      );
    }
    onProgress?.call(0.20);

    Float32List? stageLatent;
    for (var stage = 0; stage < 4; stage++) {
      _checkCancel(shouldCancel);
      final params = await network.runStage(stage, base);
      final scales = masks.applyMask(params.scalesSupp, stage);
      final means = masks.applyMask(params.meansSupp, stage);
      final indexes = aeicBuildIndexes(masks.squeeze(scales));
      final symbols = decoder.decodeStream(indexes, kAeicYCdfGroup);
      final meansSqueezed = masks.squeeze(means);
      if (symbols.length != meansSqueezed.length) {
        throw StateError(
          'stage $stage decoded ${symbols.length} symbols but the context '
          'model produced ${meansSqueezed.length} means',
        );
      }
      final latentSqueezed = Float32List(symbols.length);
      for (var i = 0; i < symbols.length; i++) {
        latentSqueezed[i] = symbols[i] + meansSqueezed[i];
      }
      stageLatent = masks.unsqueeze(latentSqueezed, stage);
      if (stage < 3) {
        base = masks.mergeContext(base, stageLatent, stage);
      }
      onProgress?.call(0.20 + 0.19 * (stage + 1));
    }

    // y_hat = base * (1 - mask_3) + y_hat_3
    return masks.mergeContext(base, stageLatent!, 3);
  }

  static void _checkCancel(AeicCancelCheck? shouldCancel) {
    if (shouldCancel?.call() == true) {
      throw const AeicEntropyCancelled();
    }
  }
}

/// Cancellation signal raised out of the entropy loop at a stage boundary.
class AeicEntropyCancelled implements Exception {
  const AeicEntropyCancelled();

  @override
  String toString() => 'Image processing stopped.';
}