core/dwt.ts

import type { Buffer2D, DwtResult } from './types.js';

/**
 * Create a 2D buffer
 */
export function createBuffer2D(width: number, height: number): Buffer2D {
  return { data: new Float64Array(width * height), width, height };
}

/**
 * Copy a Buffer2D
 */
export function copyBuffer2D(buf: Buffer2D): Buffer2D {
  return { data: new Float64Array(buf.data), width: buf.width, height: buf.height };
}

/**
 * Get value from 2D buffer
 */
function get(buf: Buffer2D, x: number, y: number): number {
  return buf.data[y * buf.width + x];
}

/**
 * Set value in 2D buffer
 */
function set(buf: Buffer2D, x: number, y: number, val: number): void {
  buf.data[y * buf.width + x] = val;
}

// Reusable temp buffer for Haar transforms — grown as needed
let _haarTmp: Float64Array | null = null;

function getHaarTmp(len: number): Float64Array {
  if (!_haarTmp || _haarTmp.length < len) {
    _haarTmp = new Float64Array(len);
  }
  return _haarTmp;
}

/**
 * 1D Haar forward transform (in-place on a row/column slice)
 */
function haarForward1D(input: Float64Array, length: number): void {
  const half = length >> 1;
  const tmp = getHaarTmp(length);
  const s = Math.SQRT1_2; // 1/sqrt(2)
  for (let i = 0; i < half; i++) {
    const a = input[2 * i];
    const b = input[2 * i + 1];
    tmp[i] = (a + b) * s;        // approximation (low)
    tmp[half + i] = (a - b) * s; // detail (high)
  }
  for (let i = 0; i < length; i++) input[i] = tmp[i];
}

/**
 * 1D Haar inverse transform (in-place on a row/column slice)
 */
function haarInverse1D(input: Float64Array, length: number): void {
  const half = length >> 1;
  const tmp = getHaarTmp(length);
  const s = Math.SQRT1_2;
  for (let i = 0; i < half; i++) {
    const low = input[i];
    const high = input[half + i];
    tmp[2 * i] = (low + high) * s;
    tmp[2 * i + 1] = (low - high) * s;
  }
  for (let i = 0; i < length; i++) input[i] = tmp[i];
}

/**
 * Forward 2D Haar DWT (one level)
 * Transforms the top-left (w x h) region of the buffer
 */
function dwtForward2DLevel(buf: Buffer2D, w: number, h: number): void {
  // Transform rows
  const rowBuf = new Float64Array(w);
  for (let y = 0; y < h; y++) {
    const off = y * buf.width;
    for (let x = 0; x < w; x++) rowBuf[x] = buf.data[off + x];
    haarForward1D(rowBuf, w);
    for (let x = 0; x < w; x++) buf.data[off + x] = rowBuf[x];
  }
  // Transform columns
  const colBuf = new Float64Array(h);
  for (let x = 0; x < w; x++) {
    for (let y = 0; y < h; y++) colBuf[y] = buf.data[y * buf.width + x];
    haarForward1D(colBuf, h);
    for (let y = 0; y < h; y++) buf.data[y * buf.width + x] = colBuf[y];
  }
}

/**
 * Inverse 2D Haar DWT (one level)
 */
function dwtInverse2DLevel(buf: Buffer2D, w: number, h: number): void {
  // Inverse columns
  const colBuf = new Float64Array(h);
  for (let x = 0; x < w; x++) {
    for (let y = 0; y < h; y++) colBuf[y] = buf.data[y * buf.width + x];
    haarInverse1D(colBuf, h);
    for (let y = 0; y < h; y++) buf.data[y * buf.width + x] = colBuf[y];
  }
  // Inverse rows
  const rowBuf = new Float64Array(w);
  for (let y = 0; y < h; y++) {
    const off = y * buf.width;
    for (let x = 0; x < w; x++) rowBuf[x] = buf.data[off + x];
    haarInverse1D(rowBuf, w);
    for (let x = 0; x < w; x++) buf.data[off + x] = rowBuf[x];
  }
}

/**
 * Extract a subband from the DWT buffer
 * After one level of DWT on a (w x h) region:
 *   LL = top-left (w/2 x h/2)
 *   LH = top-right (w/2 x h/2) — vertical detail
 *   HL = bottom-left (w/2 x h/2) — horizontal detail
 *   HH = bottom-right (w/2 x h/2) — diagonal detail
 */
export function extractSubband(
  buf: Buffer2D,
  w: number,
  h: number,
  band: 'LL' | 'LH' | 'HL' | 'HH'
): Buffer2D {
  const hw = w >> 1;
  const hh = h >> 1;
  const offX = band === 'LH' || band === 'HH' ? hw : 0;
  const offY = band === 'HL' || band === 'HH' ? hh : 0;
  const out = createBuffer2D(hw, hh);
  for (let y = 0; y < hh; y++) {
    const srcOff = (offY + y) * buf.width + offX;
    const dstOff = y * hw;
    for (let x = 0; x < hw; x++) {
      out.data[dstOff + x] = buf.data[srcOff + x];
    }
  }
  return out;
}

/**
 * Write a subband back into the DWT buffer
 */
export function writeSubband(
  buf: Buffer2D,
  w: number,
  h: number,
  band: 'LL' | 'LH' | 'HL' | 'HH',
  subband: Buffer2D
): void {
  const hw = w >> 1;
  const hh = h >> 1;
  const offX = band === 'LH' || band === 'HH' ? hw : 0;
  const offY = band === 'HL' || band === 'HH' ? hh : 0;
  for (let y = 0; y < hh; y++) {
    const dstOff = (offY + y) * buf.width + offX;
    const srcOff = y * hw;
    for (let x = 0; x < hw; x++) {
      buf.data[dstOff + x] = subband.data[srcOff + x];
    }
  }
}

/**
 * Multi-level forward DWT
 * Returns the modified buffer and subband info for reconstruction
 */
export function dwtForward(input: Buffer2D, levels: number): { buf: Buffer2D; dims: Array<{ w: number; h: number }> } {
  const buf = copyBuffer2D(input);
  const dims: Array<{ w: number; h: number }> = [];
  let w = buf.width;
  let h = buf.height;

  for (let l = 0; l < levels; l++) {
    // Ensure even dimensions
    w = w & ~1;
    h = h & ~1;
    dims.push({ w, h });
    dwtForward2DLevel(buf, w, h);
    w >>= 1;
    h >>= 1;
  }

  return { buf, dims };
}

/**
 * Multi-level inverse DWT
 */
export function dwtInverse(buf: Buffer2D, dims: Array<{ w: number; h: number }>): void {
  for (let l = dims.length - 1; l >= 0; l--) {
    const { w, h } = dims[l];
    dwtInverse2DLevel(buf, w, h);
  }
}

/**
 * Convert a Uint8Array Y plane to a Float64 Buffer2D
 */
export function yPlaneToBuffer(yPlane: Uint8Array, width: number, height: number): Buffer2D {
  const buf = createBuffer2D(width, height);
  for (let i = 0; i < width * height; i++) {
    buf.data[i] = yPlane[i];
  }
  return buf;
}

/**
 * Convert a Float64 Buffer2D back to a Uint8Array Y plane (clamped to 0–255)
 */
export function bufferToYPlane(buf: Buffer2D): Uint8Array {
  const out = new Uint8Array(buf.width * buf.height);
  for (let i = 0; i < out.length; i++) {
    out[i] = Math.max(0, Math.min(255, Math.round(buf.data[i])));
  }
  return out;
}