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

import 'package:flutter_test/flutter_test.dart';
import 'package:meshcore_open/services/entropy_tables.dart';
import 'package:meshcore_open/services/image_codec_backend.dart'
    show AeicRansCoders;
import 'package:meshcore_open/services/image_codec_entropy.dart';

/// Golden vectors for the entropy layer — the arithmetic between the ONNX
/// tensors and the rANS coder.
///
/// GENERATED BY: `aic/exp/export_entropy_layer_golden.py`, which runs the real
/// `torch` ops from `aic/aeic/src/codec/codec_practical.py`
/// (`get_mask_four_parts`, `sequeeze`, `torch.round`, `my_build_indexes`) on
/// inputs built from a closed-form integer recipe. Every constant in that
/// recipe is a power-of-two fraction or a plain float64 division, so Dart and
/// numpy reproduce the float32 inputs bit-for-bit and the comparison is real
/// rather than a re-implementation checking itself.
///
/// Regenerate with:
///     cd /Users/Zach/Documents/mycode/aic
///     AEIC_DEVICE=cpu .venv/bin/python exp/export_entropy_layer_golden.py
///
/// WHY THIS MATTERS: none of these failures are loud. A wrong mask permutation,
/// a `sequeeze` that folds the wrong channels, or a rounding tie resolved away
/// from zero instead of to even does not throw — it desynchronises rANS and
/// produces a sharp, plausible, wrong image.
class _Golden {
  // --- small case: C = 8, H = 4, W = 4, squeezed length 32 ---

  /// `mask_i` flattened over `[1, 8, 4, 4]`, '1' where live. Straight from
  /// `get_mask_four_parts(1, 8, 4, 4)`.
  static const List<String> masks = <String>[
    '10100000101000001010000010100000010100000101000001010000010100000000101000001010000010100000101000000101000001010000010100000101',
    '00000101000001010000010100000101000010100000101000001010000010100101000001010000010100000101000010100000101000001010000010100000',
    '00001010000010100000101000001010000001010000010100000101000001011010000010100000101000001010000001010000010100000101000001010000',
    '01010000010100000101000001010000101000001010000010100000101000000000010100000101000001010000010100001010000010100000101000001010',
  ];

  static const List<List<int>> symbols = <List<int>>[
    <int>[
      -1,
      -1,
      -2,
      4,
      0,
      0,
      6,
      -3,
      2,
      1,
      1,
      -1,
      -5,
      1,
      0,
      0,
      -1,
      -1,
      -2,
      4,
      0,
      0,
      5,
      -3,
      2,
      1,
      1,
      -1,
      -5,
      1,
      0,
      0,
    ],
    <int>[
      -1,
      -1,
      4,
      4,
      0,
      0,
      0,
      5,
      2,
      1,
      -1,
      -1,
      3,
      -5,
      1,
      0,
      -1,
      -2,
      4,
      4,
      0,
      0,
      0,
      5,
      1,
      1,
      -1,
      -1,
      3,
      -5,
      1,
      0,
    ],
    <int>[
      -1,
      -2,
      -2,
      4,
      1,
      0,
      0,
      5,
      2,
      0,
      -1,
      -1,
      3,
      -5,
      1,
      0,
      -1,
      -2,
      -2,
      4,
      0,
      0,
      0,
      5,
      2,
      -1,
      -1,
      -1,
      3,
      -5,
      1,
      0,
    ],
    <int>[
      -1,
      -1,
      -2,
      -2,
      0,
      0,
      5,
      -3,
      2,
      1,
      -1,
      -1,
      -5,
      -5,
      0,
      0,
      -1,
      -1,
      -2,
      -2,
      0,
      0,
      5,
      -3,
      2,
      1,
      -1,
      -1,
      -5,
      1,
      0,
      0,
    ],
  ];

  static const List<List<int>> indexes = <List<int>>[
    <int>[
      -1,
      6,
      -1,
      8,
      12,
      15,
      13,
      16,
      0,
      9,
      3,
      10,
      13,
      16,
      14,
      17,
      0,
      9,
      3,
      10,
      13,
      16,
      14,
      17,
      5,
      11,
      7,
      12,
      15,
      17,
      15,
      18,
    ],
    <int>[
      14,
      11,
      15,
      12,
      7,
      0,
      8,
      2,
      15,
      13,
      16,
      14,
      9,
      4,
      11,
      6,
      15,
      13,
      16,
      14,
      10,
      4,
      11,
      6,
      16,
      14,
      17,
      15,
      11,
      7,
      12,
      9,
    ],
    <int>[
      11,
      14,
      12,
      15,
      -1,
      8,
      0,
      9,
      12,
      16,
      13,
      16,
      2,
      10,
      5,
      11,
      13,
      16,
      13,
      16,
      3,
      10,
      5,
      11,
      14,
      17,
      15,
      17,
      7,
      12,
      8,
      13,
    ],
    <int>[
      5,
      -1,
      7,
      0,
      15,
      12,
      15,
      13,
      8,
      1,
      10,
      4,
      16,
      14,
      16,
      14,
      8,
      2,
      10,
      4,
      16,
      14,
      16,
      14,
      11,
      6,
      12,
      8,
      17,
      15,
      17,
      16,
    ],
  ];

  // --- full case: C = 256, H = W = 16 (the shipping geometry) ---
  //
  // 16,384 symbols per stage is too much to embed, so the golden is a
  // fingerprint: total, extremes, and the first and last twelve values. Any
  // permutation error moves the head or the tail; any arithmetic error moves
  // the sum.
  static const List<int> fullSymbolSum = <int>[-17, 40, -10, 29];
  static const List<int> fullIndexSum = <int>[412240, 412285, 412231, 412219];
  static const List<List<int>> fullSymbolHead = <List<int>>[
    <int>[-1, -4, -2, 1, 4, -1, -5, -2, 0, 3, -1, -5],
    <int>[-3, 0, 3, 5, 0, -4, -1, 2, 4, 1, -4, -2],
    <int>[0, 3, -1, 2, -3, 0, 2, -1, 1, -4, -1, 2],
    <int>[-3, -1, 2, -2, 1, 3, -2, 1, 4, 0, 3, 5],
  ];
  static const List<List<int>> fullSymbolTail = <List<int>>[
    <int>[-4, -2, 1, 4, -2, 1, -3, 0, 3, -2, -6, -3],
    <int>[-2, 0, -3, -1, 2, 5, 0, -4, -1, 1, 4, 6],
    <int>[-5, -2, 0, 3, -1, -6, -3, 0, 2, 5, 1, -4],
    <int>[-1, 1, 4, 0, -4, -2, 1, 3, 6, 2, -2, 0],
  ];
  static const List<List<int>> fullIndexHead = <List<int>>[
    <int>[-1, -1, -1, 1, 0, 4, 4, 6, 6, 8, 8, 9],
    <int>[0, 0, 4, 3, 6, 6, 8, 7, 9, 9, 10, 10],
    <int>[0, 2, 2, 5, 4, 7, 7, 8, 8, 10, 10, 11],
    <int>[-1, -1, 0, 0, 3, 2, 5, 5, 7, 7, 9, 8],
  ];
  static const List<List<int>> fullIndexTail = <List<int>>[
    <int>[16, 17, 17, 18, 17, 18, 18, 18, 18, 19, 19, 19],
    <int>[16, 16, 17, 16, 17, 17, 18, 18, 18, 18, 19, 18],
    <int>[16, 16, 16, 17, 17, 17, 17, 18, 18, 18, 18, 19],
    <int>[17, 17, 17, 17, 18, 18, 18, 18, 19, 19, 19, 19],
  ];

  /// `my_build_indexes` on a spread of scales, isolated from everything else.
  static const List<double> probeScales = <double>[
    0.0,
    9.999999974752427e-07,
    0.05000000074505806,
    0.07989999651908875,
    0.07999999821186066,
    0.10999999940395355,
    0.11000010371208191,
    0.5,
    1.0,
    2.0,
    7.5,
    63.900001525878906,
    255.89999389648438,
    256.0,
    10000.0,
  ];
  static const List<int> probeIndexes = <int>[
    -1,
    -1,
    -1,
    -1,
    0,
    0,
    0,
    12,
    17,
    23,
    34,
    51,
    62,
    63,
    63,
  ];
}

/// The same closed-form recipe the generator uses, in Dart.
///
/// `y`, `means_supp` and `scales_supp` for a `[1, C, H, W]` tensor. Kept
/// byte-identical to `recipe()` in `export_entropy_layer_golden.py`.
({Float32List y, Float32List means, Float32List scales}) _recipe(
  int channels,
  int height,
  int width,
) {
  final n = channels * height * width;
  final y = Float32List(n);
  final means = Float32List(n);
  final scales = Float32List(n);
  var i = 0;
  for (var c = 0; c < channels; c++) {
    for (var h = 0; h < height; h++) {
      for (var w = 0; w < width; w++, i++) {
        y[i] = ((c * 3 + h * 17 + w * 11) % 61 - 30) * 0.125;
        means[i] = ((c * 7 + h * 13 + w * 29) % 97 - 48) * 0.0625;
        scales[i] = ((c * 11 + h * 5 + w * 3) % 700) * 0.01 + 0.001;
      }
    }
  }
  return (y: y, means: means, scales: scales);
}

/// One stage of `compress()`: mask, quantize, fold, index.
({Int16List symbols, Int16List indexes}) _runStage(
  AeicMaskSet masks,
  Float32List y,
  Float32List meansSupp,
  Float32List scalesSupp,
  int stage,
) {
  final means = masks.applyMask(meansSupp, stage);
  final scales = masks.applyMask(scalesSupp, stage);
  final maskedY = masks.applyMask(y, stage);
  final yq = Float32List(y.length);
  for (var i = 0; i < y.length; i++) {
    yq[i] = roundHalfToEven(f32(maskedY[i] - means[i]));
  }
  return (
    symbols: aeicToSymbols(masks.squeeze(yq)),
    indexes: aeicBuildIndexes(masks.squeeze(scales)),
  );
}

void main() {
  group('AeicEntropyGeometry', () {
    test('512x512 ft32 matches the shapes the bitstream format assumes', () {
      final g = AeicEntropyGeometry.forResolution(512);
      expect(g.yShape, <int>[1, 256, 16, 16]);
      expect(g.zShape, <int>[1, 128, 4, 4]);
      expect(g.squeezedChannels, 64);
      // From aic/results/rans_port_spec.md §1.
      expect(g.zElements, 2048);
      expect(g.symbolsPerStage, 16384);
      expect(g.totalEntries, 67584);
    });

    test('z is ceil(y/4), not floor', () {
      // 320/32 = 10 -> z must be 3, matching compress()'s reflect padding.
      final g = AeicEntropyGeometry.forResolution(320);
      expect(g.yHeight, 10);
      expect(g.zHeight, 3);
    });

    test('rejects a resolution g_a cannot downsample by 32', () {
      expect(() => AeicEntropyGeometry.forResolution(500), throwsArgumentError);
      expect(() => AeicEntropyGeometry.forResolution(0), throwsArgumentError);
    });
  });

  group('AeicMaskSet', () {
    test('the four masks are get_mask_four_parts, element for element', () {
      final masks = AeicMaskSet(
        AeicEntropyGeometry.forResolution(128, yChannels: 8),
      );
      // 128/32 = 4, so this is exactly the C=8 H=W=4 case in the golden.
      for (var stage = 0; stage < 4; stage++) {
        final tensor = masks.maskTensor(stage);
        final actual = tensor.map((v) => v == 1.0 ? '1' : '0').join();
        expect(actual, _Golden.masks[stage], reason: 'mask_$stage');
      }
    });

    test('every position is claimed by exactly one stage, per channel', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final masks = AeicMaskSet(geometry);
      final counts = Uint8List(geometry.yElements);
      for (var stage = 0; stage < 4; stage++) {
        final tensor = masks.maskTensor(stage);
        for (var i = 0; i < tensor.length; i++) {
          counts[i] += tensor[i].toInt();
        }
      }
      expect(counts.every((c) => c == 1), isTrue);
    });

    test('each mask carries exactly a quarter of the tensor', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final masks = AeicMaskSet(geometry);
      for (var stage = 0; stage < 4; stage++) {
        final live = masks.maskTensor(stage).where((v) => v == 1.0).length;
        expect(live, geometry.symbolsPerStage);
      }
    });

    test('squeeze then unsqueeze is the identity on a masked tensor', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final masks = AeicMaskSet(geometry);
      final recipe = _recipe(256, 16, 16);
      for (var stage = 0; stage < 4; stage++) {
        final masked = masks.applyMask(recipe.y, stage);
        final restored = masks.unsqueeze(masks.squeeze(masked), stage);
        expect(restored, masked, reason: 'stage $stage');
      }
    });

    test('mergeContext replaces exactly the stage mask', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final masks = AeicMaskSet(geometry);
      final base = _recipe(256, 16, 16).y;
      final stageLatent = masks.applyMask(_recipe(256, 16, 16).means, 2);
      final merged = masks.mergeContext(base, stageLatent, 2);
      final mask = masks.maskTensor(2);
      for (var i = 0; i < merged.length; i++) {
        expect(merged[i], mask[i] == 1.0 ? stageLatent[i] : base[i]);
      }
    });
  });

  group('roundHalfToEven', () {
    test('ties go to even, unlike Dart round()', () {
      expect(roundHalfToEven(0.5), 0.0);
      expect(roundHalfToEven(1.5), 2.0);
      expect(roundHalfToEven(2.5), 2.0);
      expect(roundHalfToEven(-0.5), 0.0);
      expect(roundHalfToEven(-1.5), -2.0);
      expect(roundHalfToEven(-2.5), -2.0);
      // Dart disagrees on every one of those ties.
      expect((-0.5).roundToDouble(), -1.0);
    });

    test('non-ties are ordinary rounding', () {
      expect(roundHalfToEven(0.49), 0.0);
      expect(roundHalfToEven(0.51), 1.0);
      expect(roundHalfToEven(-1.51), -2.0);
      expect(roundHalfToEven(-1.49), -1.0);
      expect(roundHalfToEven(7.0), 7.0);
    });
  });

  group('aeicBuildIndexes', () {
    test('matches my_build_indexes on the probe scales', () {
      final scales = Float32List.fromList(_Golden.probeScales);
      expect(aeicBuildIndexes(scales), _Golden.probeIndexes);
    });

    test('the 0.08 threshold is strict, and clamps do not leak', () {
      // 0.08 itself is NOT skipped; it clamps to row 0 because ln(0.08) is
      // below ln(0.11).
      expect(aeicBuildIndexes(Float32List.fromList(<double>[0.08])).first, 0);
      expect(
        aeicBuildIndexes(Float32List.fromList(<double>[0.0799])).first,
        -1,
      );
      expect(
        aeicBuildIndexes(Float32List.fromList(<double>[1e9])).first,
        kAeicScalesLevels - 1,
      );
    });
  });

  group('aeicZIndexes', () {
    test('is the channel arange broadcast over H*W', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final indexes = aeicZIndexes(geometry);
      expect(indexes.length, 2048);
      expect(indexes.first, 0);
      expect(indexes[15], 0);
      expect(indexes[16], 1);
      expect(indexes.last, 127);
      // Verified against results/golden/vectors/kodim01.gv, which stores the
      // exact int16 array the C++ coder was given.
      for (var i = 0; i < indexes.length; i++) {
        expect(indexes[i], i ~/ 16);
      }
    });
  });

  group('aeicToSymbols', () {
    test('rejects a value the int16 wire format cannot carry', () {
      expect(
        () => aeicToSymbols(Float32List.fromList(<double>[40000.0])),
        throwsStateError,
      );
      expect(
        () => aeicToSymbols(Float32List.fromList(<double>[1.5])),
        throwsStateError,
      );
      expect(aeicToSymbols(Float32List.fromList(<double>[-3.0])), <int>[-3]);
    });
  });

  group('aeicRgbToChw', () {
    test('reproduces ToTensor + Normalize([0.5], [0.5])', () {
      final chw = aeicRgbToChw(Uint8List.fromList(<int>[0, 128, 255]), 1);
      expect(chw.length, 3);
      expect(chw[0], -1.0);
      expect(chw[1], closeTo(0.00392, 1e-4));
      expect(chw[2], 1.0);
    });

    test('is channel-planar, not interleaved', () {
      final rgb = Uint8List(4 * 3);
      for (var i = 0; i < 4; i++) {
        rgb[i * 3] = 255; // R
        rgb[i * 3 + 1] = 0; // G
        rgb[i * 3 + 2] = 0; // B
      }
      final chw = aeicRgbToChw(rgb, 2);
      expect(chw.sublist(0, 4), <double>[1.0, 1.0, 1.0, 1.0]);
      expect(chw.sublist(4, 8), <double>[-1.0, -1.0, -1.0, -1.0]);
    });

    test('rejects a byte count that is not the stated square', () {
      expect(() => aeicRgbToChw(Uint8List(11), 2), throwsArgumentError);
    });
  });

  group('four-stage symbol packing (golden)', () {
    test('C=8 H=4 W=4: symbols and indexes match torch exactly', () {
      final geometry = AeicEntropyGeometry.forResolution(128, yChannels: 8);
      final masks = AeicMaskSet(geometry);
      final recipe = _recipe(8, 4, 4);
      for (var stage = 0; stage < 4; stage++) {
        final out = _runStage(
          masks,
          recipe.y,
          recipe.means,
          recipe.scales,
          stage,
        );
        expect(out.symbols, _Golden.symbols[stage], reason: 'symbols $stage');
        expect(out.indexes, _Golden.indexes[stage], reason: 'indexes $stage');
      }
    });

    test('C=256 H=W=16 (shipping geometry): fingerprint matches torch', () {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final masks = AeicMaskSet(geometry);
      final recipe = _recipe(256, 16, 16);
      for (var stage = 0; stage < 4; stage++) {
        final out = _runStage(
          masks,
          recipe.y,
          recipe.means,
          recipe.scales,
          stage,
        );
        expect(out.symbols.length, 16384);
        expect(
          out.symbols.fold<int>(0, (a, b) => a + b),
          _Golden.fullSymbolSum[stage],
          reason: 'symbol sum $stage',
        );
        expect(out.symbols.sublist(0, 12), _Golden.fullSymbolHead[stage]);
        expect(out.symbols.sublist(16372), _Golden.fullSymbolTail[stage]);
        expect(
          out.indexes.fold<int>(0, (a, b) => a + b),
          _Golden.fullIndexSum[stage],
          reason: 'index sum $stage',
        );
        expect(out.indexes.sublist(0, 12), _Golden.fullIndexHead[stage]);
        expect(out.indexes.sublist(16372), _Golden.fullIndexTail[stage]);
      }
    });
  });

  group('AeicEntropyCodec', () {
    test(
      'encode pushes z, y0, y1, y2, y3 in that order and no other',
      () async {
        final geometry = AeicEntropyGeometry.forResolution(512);
        final network = _FakeNetwork(geometry);
        final coders = _RecordingCoders();
        final codec = AeicEntropyCodec(
          geometry: geometry,
          network: network,
          coders: coders,
        );
        final progress = <double>[];
        final stream = await codec.encode(
          Uint8List(512 * 512 * 3),
          onProgress: progress.add,
        );

        expect(coders.encoder.groups, <int>[
          kAeicZCdfGroup,
          kAeicYCdfGroup,
          kAeicYCdfGroup,
          kAeicYCdfGroup,
          kAeicYCdfGroup,
        ]);
        expect(coders.encoder.lengths, <int>[2048, 16384, 16384, 16384, 16384]);
        expect(stream, isNotEmpty);
        expect(progress.last, 1.0);
        // The image really was normalized before the graph saw it: 0 -> -1.
        expect(network.lastInput!.first, -1.0);
      },
    );

    test('encode refuses a graph that returns the wrong z size', () async {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final codec = AeicEntropyCodec(
        geometry: geometry,
        network: _FakeNetwork(geometry, zElements: 7),
        coders: _RecordingCoders(),
      );
      await expectLater(
        codec.encode(Uint8List(512 * 512 * 3)),
        throwsStateError,
      );
    });

    test('encode honours shouldCancel between stages', () async {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final codec = AeicEntropyCodec(
        geometry: geometry,
        network: _FakeNetwork(geometry),
        coders: _RecordingCoders(),
      );
      await expectLater(
        codec.encode(Uint8List(512 * 512 * 3), shouldCancel: () => true),
        throwsA(isA<AeicEntropyCancelled>()),
      );
    });

    test('decode reports a send-side-only graph instead of guessing', () async {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final codec = AeicEntropyCodec(
        geometry: geometry,
        network: _FakeNetwork(geometry, decodeSide: false),
        coders: _RecordingCoders(),
      );
      await expectLater(
        codec.decodeToLatent(Uint8List(16)),
        throwsA(isA<AeicEntropyUnavailable>()),
      );
    });

    test('decode walks the stages in order and rebuilds y_hat', () async {
      final geometry = AeicEntropyGeometry.forResolution(512);
      final network = _FakeNetwork(geometry);
      final coders = _RecordingCoders();
      final codec = AeicEntropyCodec(
        geometry: geometry,
        network: network,
        coders: coders,
      );
      final yHat = await codec.decodeToLatent(Uint8List(16));

      expect(network.stageCalls, <int>[0, 1, 2, 3]);
      expect(coders.decoder.groups, <int>[
        kAeicZCdfGroup,
        kAeicYCdfGroup,
        kAeicYCdfGroup,
        kAeicYCdfGroup,
        kAeicYCdfGroup,
      ]);
      expect(yHat.length, geometry.yElements);
      // The fake decoder returns symbol 1 everywhere and the fake network
      // returns means 0, so every position of y_hat must be exactly 1 — which
      // only holds if the four masks tile the tensor and each stage's squeezed
      // symbols were unsqueezed back into the right channel group.
      expect(yHat.every((v) => v == 1.0), isTrue);
    });
  });
  // The end-to-end proof: the entropy layer plus the pure-Dart range coder
  // reproduce, byte for byte, the bitstreams the C++ coder produced for real
  // images — and decode them back to the same symbols.
  group('bitstream round trip against the C++ golden vectors', () {
    final Directory goldenDir = _resolveGoldenDir();
    final EntropyTables tables = EntropyTables.parse(
      File('${goldenDir.path}/aeic_cdf_ft32.bin').readAsBytesSync(),
    );
    final geometry = AeicEntropyGeometry.forResolution(512);
    final masks = AeicMaskSet(geometry);

    // kodim01 is the plain case; kodim23 and image2 both carry index -1
    // (scales below 0.08), which the coder must skip on encode and read back as
    // a literal 0.
    for (final name in <String>['kodim01', 'kodim23', 'image2']) {
      test('$name: symbols -> bitstream -> symbols', () async {
        final vector = _readGoldenVector(
          File('${goldenDir.path}/vectors/$name.gv').readAsBytesSync(),
        );
        final expectedStream = File(
          '${goldenDir.path}/vectors/$name.bin',
        ).readAsBytesSync();

        // Rebuild the pre-fold tensors the graph would have produced. `squeeze`
        // is a bijection on masked tensors, so unsqueezing the golden arrays
        // recovers a legitimate input and the loop's own fold has to invert it.
        final yQ = <Float32List>[];
        final scales = <Float32List>[];
        for (var stage = 0; stage < 4; stage++) {
          final symbols = vector['y_q$stage']!;
          final indexes = vector['y_indexes$stage']!;
          final squeezedY = Float32List(symbols.length);
          final squeezedS = Float32List(indexes.length);
          for (var i = 0; i < symbols.length; i++) {
            squeezedY[i] = symbols[i].toDouble();
            squeezedS[i] = _scaleForIndex(indexes[i]);
          }
          // Self-check: the synthesized scales must land back on the exact
          // golden indexes, or this test is measuring the wrong thing.
          expect(aeicBuildIndexes(squeezedS), indexes, reason: 'stage $stage');
          yQ.add(masks.unsqueeze(squeezedY, stage));
          scales.add(masks.unsqueeze(squeezedS, stage));
        }
        final zQ = Float32List(vector['z_q']!.length);
        for (var i = 0; i < zQ.length; i++) {
          zQ[i] = vector['z_q']![i].toDouble();
        }

        final network = _ReplayNetwork(
          geometry: geometry,
          tensors: AeicEncodeSideTensors(zQ: zQ, yQ: yQ, scales: scales),
        );
        final codec = AeicEntropyCodec(
          geometry: geometry,
          network: network,
          coders: AeicRansCoders(tables),
        );

        // ENCODE: byte-for-byte against the C++ coder's output.
        final stream = await codec.encode(Uint8List(512 * 512 * 3));
        expect(stream, expectedStream, reason: 'bitstream for $name');

        // DECODE: the same bytes back to the same symbols. means are zero, so
        // y_hat is exactly the four stages' symbols tiled back into place —
        // which only holds if every mask, fold and unfold agrees with encode.
        final yHat = await codec.decodeToLatent(stream);
        var expected = Float32List(geometry.yElements);
        for (var stage = 0; stage < 4; stage++) {
          expected = masks.mergeContext(expected, yQ[stage], stage);
        }
        expect(yHat, expected, reason: 'y_hat for $name');
        expect(network.stageCalls, <int>[0, 1, 2, 3]);
      });
    }
  });
}

Directory _resolveGoldenDir() {
  for (final candidate in <String>[
    'test/services/golden',
    '../test/services/golden',
    'golden',
  ]) {
    final dir = Directory(candidate);
    if (dir.existsSync()) {
      return dir;
    }
  }
  throw StateError(
    'golden vectors not found; expected test/services/golden relative to the '
    'package root',
  );
}

/// Parses the `.gv` container documented in `aic/results/rans_port_spec.md` §10.
Map<String, Int16List> _readGoldenVector(Uint8List bytes) {
  final data = ByteData.sublistView(bytes);
  const magic = <int>[0x41, 0x45, 0x49, 0x43, 0x47, 0x56, 0x00, 0x01];
  for (var i = 0; i < magic.length; i++) {
    if (bytes[i] != magic[i]) {
      throw StateError('not a .gv container');
    }
  }
  final count = data.getUint32(12, Endian.little);
  final names = <String>[];
  final counts = <int>[];
  final dtypes = <int>[];
  var off = 16;
  for (var i = 0; i < count; i++) {
    final raw = bytes.sublist(off, off + 16);
    final end = raw.indexOf(0);
    names.add(String.fromCharCodes(raw.sublist(0, end < 0 ? 16 : end)));
    dtypes.add(data.getUint32(off + 16, Endian.little));
    counts.add(data.getUint32(off + 20, Endian.little));
    off += 24;
  }
  final out = <String, Int16List>{};
  for (var i = 0; i < count; i++) {
    if (dtypes[i] != 0) {
      throw StateError('${names[i]} is not int16');
    }
    final values = Int16List(counts[i]);
    for (var j = 0; j < counts[i]; j++) {
      values[j] = data.getInt16(off + j * 2, Endian.little);
    }
    off += counts[i] * 2;
    out[names[i]] = values;
  }
  return out;
}

/// A scale that `my_build_indexes` maps back to exactly [index].
///
/// Row centres, so float32 rounding cannot push one over a boundary; index -1
/// means "skipped", which any scale below 0.08 produces.
double _scaleForIndex(int index) {
  if (index < 0) {
    return 0.0;
  }
  return f32(math.exp(kAeicLogScaleMin + (index + 0.5) * kAeicLogScaleStep));
}

/// Replays fixed tensors as if they came from the graph, for both directions.
class _ReplayNetwork implements AeicEntropyNetwork {
  final AeicEntropyGeometry geometry;
  final AeicEncodeSideTensors tensors;
  final List<int> stageCalls = <int>[];

  _ReplayNetwork({required this.geometry, required this.tensors});

  @override
  bool get supportsDecodeSide => true;

  @override
  Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw) async =>
      tensors;

  @override
  Future<Float32List> runHyperSynthesis(Float32List zQ) async {
    // The real h_s consumes z_hat; here the only thing under test is that the
    // decoded z symbols reach it. Assert that and hand back a zero context.
    expect(zQ.length, geometry.zElements);
    for (var i = 0; i < zQ.length; i++) {
      expect(zQ[i], tensors.zQ[i], reason: 'z symbol $i');
    }
    return Float32List(geometry.yElements);
  }

  @override
  Future<AeicStageParams> runStage(int stage, Float32List base) async {
    stageCalls.add(stage);
    return AeicStageParams(
      meansSupp: Float32List(geometry.yElements),
      scalesSupp: tensors.scales[stage],
    );
  }
}

/// A stand-in for the ONNX graph: shapes and call order, no arithmetic.
class _FakeNetwork implements AeicEntropyNetwork {
  final AeicEntropyGeometry geometry;
  final int? zElements;
  final bool decodeSide;
  final List<int> stageCalls = <int>[];
  Float32List? lastInput;

  _FakeNetwork(this.geometry, {this.zElements, this.decodeSide = true});

  @override
  bool get supportsDecodeSide => decodeSide;

  @override
  Future<AeicEncodeSideTensors> runEncodeSide(Float32List imageChw) async {
    lastInput = imageChw;
    return AeicEncodeSideTensors(
      zQ: Float32List(zElements ?? geometry.zElements),
      yQ: <Float32List>[
        for (var i = 0; i < 4; i++) Float32List(geometry.yElements),
      ],
      scales: <Float32List>[
        for (var i = 0; i < 4; i++)
          Float32List(geometry.yElements)
            ..fillRange(0, geometry.yElements, 1.0),
      ],
    );
  }

  @override
  Future<Float32List> runHyperSynthesis(Float32List zQ) async =>
      Float32List(geometry.yElements);

  @override
  Future<AeicStageParams> runStage(int stage, Float32List base) async {
    stageCalls.add(stage);
    return AeicStageParams(
      meansSupp: Float32List(geometry.yElements),
      scalesSupp: Float32List(geometry.yElements)
        ..fillRange(0, geometry.yElements, 1.0),
    );
  }
}

class _RecordingEncoder implements AeicRansEncoder {
  final List<int> groups = <int>[];
  final List<int> lengths = <int>[];

  @override
  void pushSymbols(Int16List symbols, Int16List indexes, int cdfGroup) {
    expect(symbols.length, indexes.length);
    groups.add(cdfGroup);
    lengths.add(symbols.length);
  }

  @override
  Uint8List finish() => Uint8List.fromList(<int>[0x11, 4, 0]);
}

class _RecordingDecoder implements AeicRansDecoder {
  final List<int> groups = <int>[];

  @override
  Int16List decodeStream(Int16List indexes, int cdfGroup) {
    groups.add(cdfGroup);
    return Int16List(indexes.length)..fillRange(0, indexes.length, 1);
  }
}

class _RecordingCoders implements AeicRansCoderFactory {
  final _RecordingEncoder encoder = _RecordingEncoder();
  final _RecordingDecoder decoder = _RecordingDecoder();

  @override
  AeicRansEncoder createEncoder() => encoder;

  @override
  AeicRansDecoder createDecoder(Uint8List bitstream) => decoder;
}
