core/embedder.ts

/**
 * High-level watermark embedder
 *
 * Takes a Y plane + payload + key + config → returns watermarked Y plane
 */

import type { WatermarkConfig, EmbedResult, Buffer2D } from './types.js';
import { yPlaneToBuffer, bufferToYPlane, dwtForward, dwtInverse, extractSubband, writeSubband } from './dwt.js';
import { dctForward8x8, dctInverse8x8, extractBlock, writeBlock, ZIGZAG_ORDER } from './dct.js';
import { dmqimEmbed } from './dmqim.js';
import { crcAppend } from './crc.js';
import { BchCodec } from './bch.js';
import { generateDithers, generatePermutation } from './keygen.js';
import { computeTileGrid, getTileOrigin, getTileBlocks } from './tiling.js';
import { blockAcEnergy, computeMaskingFactors } from './masking.js';

/**
 * Convert a 32-bit payload (4 bytes) to a bit array
 */
export function payloadToBits(payload: Uint8Array): Uint8Array {
  const bits = new Uint8Array(payload.length * 8);
  for (let i = 0; i < payload.length; i++) {
    for (let b = 0; b < 8; b++) {
      bits[i * 8 + b] = (payload[i] >> (7 - b)) & 1;
    }
  }
  return bits;
}

/**
 * Convert a bit array back to bytes
 */
export function bitsToPayload(bits: Uint8Array): Uint8Array {
  const bytes = new Uint8Array(Math.ceil(bits.length / 8));
  for (let i = 0; i < bits.length; i++) {
    if (bits[i]) {
      bytes[Math.floor(i / 8)] |= 1 << (7 - (i % 8));
    }
  }
  return bytes;
}

/**
 * Embed a watermark into a Y plane
 */
export function embedWatermark(
  yPlane: Uint8Array,
  width: number,
  height: number,
  payload: Uint8Array,
  key: string,
  config: WatermarkConfig
): EmbedResult {
  // Step 1: Encode payload
  // Always use exactly 32 bits of payload, zero-pad the BCH message to fill k
  const PAYLOAD_BITS = 32;
  let payloadBits = payloadToBits(payload);
  // Truncate or pad to exactly 32 bits
  if (payloadBits.length < PAYLOAD_BITS) {
    const padded = new Uint8Array(PAYLOAD_BITS);
    padded.set(payloadBits);
    payloadBits = padded;
  } else if (payloadBits.length > PAYLOAD_BITS) {
    payloadBits = payloadBits.subarray(0, PAYLOAD_BITS);
  }
  // CRC protects the 32-bit payload
  const withCrc = crcAppend(payloadBits, config.crc.bits);
  // Pad to fill BCH message length k
  const bchMessage = new Uint8Array(config.bch.k);
  bchMessage.set(withCrc);
  const bch = new BchCodec(config.bch);
  const encoded = bch.encode(bchMessage);
  const codedLength = encoded.length;

  // Generate interleaving permutation
  const perm = generatePermutation(key, codedLength);
  const interleaved = new Uint8Array(codedLength);
  for (let i = 0; i < codedLength; i++) {
    interleaved[perm[i]] = encoded[i];
  }

  // Step 2: Forward DWT
  const buf = yPlaneToBuffer(yPlane, width, height);
  const { buf: dwtBuf, dims } = dwtForward(buf, config.dwtLevels);

  // Step 3: Extract HL subband at deepest level
  // After N levels, the HL subband is in the bottom-left of the level-N region
  let w = dims[dims.length - 1].w;
  let h = dims[dims.length - 1].h;
  for (let l = 0; l < config.dwtLevels - 1; l++) {
    w >>= 1;
    h >>= 1;
  }
  const hlSubband = extractSubband(dwtBuf, dims[dims.length - 1].w, dims[dims.length - 1].h, 'HL');

  // Step 4: Set up tile grid
  // tilePeriod is in spatial pixels; in subband it's tilePeriod / (2^levels)
  const subbandTilePeriod = Math.floor(config.tilePeriod / (1 << config.dwtLevels));
  const tileGrid = computeTileGrid(hlSubband.width, hlSubband.height, subbandTilePeriod);

  if (tileGrid.totalTiles === 0) {
    // Frame too small for any tiles — return unchanged
    return { yPlane: new Uint8Array(yPlane), psnr: Infinity };
  }

  // Step 5: Generate dithers — same sequence reused per tile so each tile
  // is independently decodable (required for crop robustness)
  const coeffsPerBlock = config.zigzagPositions.length;
  const maxCoeffsPerTile = 1024; // Upper bound
  const dithers = generateDithers(key, maxCoeffsPerTile, config.delta);

  // Step 6: For each tile, embed the coded bits
  // Reusable block buffer to avoid allocation per block
  const blockBuf = new Float64Array(64);

  // Precompute zigzag → coefficient index mapping
  const zigCoeffIdx = new Int32Array(config.zigzagPositions.length);
  for (let z = 0; z < config.zigzagPositions.length; z++) {
    const [r, c] = ZIGZAG_ORDER[config.zigzagPositions[z]];
    zigCoeffIdx[z] = r * 8 + c;
  }

  for (let tileIdx = 0; tileIdx < tileGrid.totalTiles; tileIdx++) {
    let ditherIdx = 0; // Reset per tile — each tile uses same dither sequence
    const origin = getTileOrigin(tileGrid, tileIdx);
    const blocks = getTileBlocks(origin.x, origin.y, subbandTilePeriod, hlSubband.width, hlSubband.height);

    // Compute masking factors if enabled
    let maskingFactors: Float64Array | null = null;
    if (config.perceptualMasking && blocks.length > 0) {
      const energies = new Float64Array(blocks.length);
      for (let bi = 0; bi < blocks.length; bi++) {
        extractBlock(hlSubband.data, hlSubband.width, blocks[bi].row, blocks[bi].col, blockBuf);
        dctForward8x8(blockBuf);
        energies[bi] = blockAcEnergy(blockBuf);
      }
      maskingFactors = computeMaskingFactors(energies);
    }

    let bitIdx = 0;
    for (let bi = 0; bi < blocks.length; bi++) {
      const { row, col } = blocks[bi];
      extractBlock(hlSubband.data, hlSubband.width, row, col, blockBuf);

      // Forward DCT
      dctForward8x8(blockBuf);

      const maskFactor = maskingFactors ? maskingFactors[bi] : 1.0;
      const effectiveDelta = config.delta * maskFactor;

      // Embed bits into selected zigzag positions
      for (let z = 0; z < zigCoeffIdx.length; z++) {
        if (bitIdx >= codedLength) bitIdx = 0; // Repeat pattern

        const coeffIdx = zigCoeffIdx[z];
        const bit = interleaved[bitIdx];
        const dither = dithers[ditherIdx++];

        blockBuf[coeffIdx] = dmqimEmbed(blockBuf[coeffIdx], bit, effectiveDelta, dither);

        bitIdx++;
      }

      // Inverse DCT and write back
      dctInverse8x8(blockBuf);
      writeBlock(hlSubband.data, hlSubband.width, row, col, blockBuf);
    }
  }

  // Step 7: Write modified HL subband back and inverse DWT
  writeSubband(dwtBuf, dims[dims.length - 1].w, dims[dims.length - 1].h, 'HL', hlSubband);
  dwtInverse(dwtBuf, dims);

  // Convert back to Y plane
  const watermarkedY = bufferToYPlane(dwtBuf);

  // Compute PSNR
  let mse = 0;
  for (let i = 0; i < yPlane.length; i++) {
    const diff = yPlane[i] - watermarkedY[i];
    mse += diff * diff;
  }
  mse /= yPlane.length;
  const psnr = mse > 0 ? 10 * Math.log10(255 * 255 / mse) : Infinity;

  return { yPlane: watermarkedY, psnr };
}