/**
 * 8x8 DCT forward and inverse transforms
 * Separable implementation: rows then columns — O(N^3) instead of O(N^4)
 * Reuses temp buffers to avoid per-call allocations
 */

const N = 8;

// Precompute cosine table: cos((2i+1)*j*PI / 16) for i,j in [0,8)
const COS_TABLE = new Float64Array(N * N);
for (let i = 0; i < N; i++) {
  for (let j = 0; j < N; j++) {
    COS_TABLE[i * N + j] = Math.cos(((2 * i + 1) * j * Math.PI) / (2 * N));
  }
}

// Precompute alpha normalization table: alpha(u) * alpha(v) for all (u,v)
const ALPHA_0 = 1 / Math.sqrt(N);
const ALPHA_K = Math.sqrt(2 / N);
const ALPHA_UV = new Float64Array(N * N);
const ALPHA = new Float64Array(N);
for (let k = 0; k < N; k++) ALPHA[k] = k === 0 ? ALPHA_0 : ALPHA_K;
for (let u = 0; u < N; u++) {
  for (let v = 0; v < N; v++) {
    ALPHA_UV[u * N + v] = ALPHA[u] * ALPHA[v];
  }
}

// Precompute scaled cosine: ALPHA[k] * COS_TABLE[i*N+k] for forward transform
const SCALED_COS = new Float64Array(N * N);
for (let i = 0; i < N; i++) {
  for (let k = 0; k < N; k++) {
    SCALED_COS[i * N + k] = ALPHA[k] * COS_TABLE[i * N + k];
  }
}

// Shared temp buffers (avoids allocation per call)
const _temp64 = new Float64Array(64);
const _temp8 = new Float64Array(8);

/**
 * Forward 8x8 DCT on a block (in-place)
 * Separable: transform rows, then columns — O(8^3) = 512 ops vs O(8^4) = 4096
 */
export function dctForward8x8(block: Float64Array): void {
  // Step 1: Transform each row (spatial x → frequency v)
  for (let x = 0; x < N; x++) {
    const rowOff = x * N;
    for (let v = 0; v < N; v++) {
      let sum = 0;
      for (let y = 0; y < N; y++) {
        sum += block[rowOff + y] * COS_TABLE[y * N + v];
      }
      _temp64[rowOff + v] = ALPHA[v] * sum;
    }
  }

  // Step 2: Transform each column (spatial x → frequency u)
  for (let v = 0; v < N; v++) {
    for (let u = 0; u < N; u++) {
      let sum = 0;
      for (let x = 0; x < N; x++) {
        sum += _temp64[x * N + v] * COS_TABLE[x * N + u];
      }
      block[u * N + v] = ALPHA[u] * sum;
    }
  }
}

/**
 * Inverse 8x8 DCT on a block (in-place)
 * Separable: inverse columns, then inverse rows
 */
export function dctInverse8x8(block: Float64Array): void {
  // Step 1: Inverse transform columns (frequency u → spatial x)
  for (let v = 0; v < N; v++) {
    for (let x = 0; x < N; x++) {
      let sum = 0;
      for (let u = 0; u < N; u++) {
        sum += ALPHA[u] * block[u * N + v] * COS_TABLE[x * N + u];
      }
      _temp64[x * N + v] = sum;
    }
  }

  // Step 2: Inverse transform rows (frequency v → spatial y)
  for (let x = 0; x < N; x++) {
    const rowOff = x * N;
    for (let y = 0; y < N; y++) {
      let sum = 0;
      for (let v = 0; v < N; v++) {
        sum += ALPHA[v] * _temp64[rowOff + v] * COS_TABLE[y * N + v];
      }
      block[rowOff + y] = sum;
    }
  }
}

/**
 * Zigzag scan order for 8x8 block
 * Maps zigzag index → [row, col]
 */
export const ZIGZAG_ORDER: [number, number][] = [
  [0,0],[0,1],[1,0],[2,0],[1,1],[0,2],[0,3],[1,2],
  [2,1],[3,0],[4,0],[3,1],[2,2],[1,3],[0,4],[0,5],
  [1,4],[2,3],[3,2],[4,1],[5,0],[6,0],[5,1],[4,2],
  [3,3],[2,4],[1,5],[0,6],[0,7],[1,6],[2,5],[3,4],
  [4,3],[5,2],[6,1],[7,0],[7,1],[6,2],[5,3],[4,4],
  [3,5],[2,6],[1,7],[2,7],[3,6],[4,5],[5,4],[6,3],
  [7,2],[7,3],[6,4],[5,5],[4,6],[3,7],[4,7],[5,6],
  [6,5],[7,4],[7,5],[6,6],[5,7],[6,7],[7,6],[7,7],
];

/**
 * Extract an 8x8 block from a subband into a reusable buffer
 */
export function extractBlock(
  data: Float64Array,
  width: number,
  blockRow: number,
  blockCol: number,
  out?: Float64Array
): Float64Array {
  const block = out || new Float64Array(64);
  const startY = blockRow * 8;
  const startX = blockCol * 8;
  for (let r = 0; r < 8; r++) {
    const srcOff = (startY + r) * width + startX;
    const dstOff = r * 8;
    for (let c = 0; c < 8; c++) {
      block[dstOff + c] = data[srcOff + c];
    }
  }
  return block;
}

/**
 * Write an 8x8 block back into a subband
 */
export function writeBlock(
  data: Float64Array,
  width: number,
  blockRow: number,
  blockCol: number,
  block: Float64Array
): void {
  const startY = blockRow * 8;
  const startX = blockCol * 8;
  for (let r = 0; r < 8; r++) {
    const dstOff = (startY + r) * width + startX;
    const srcOff = r * 8;
    for (let c = 0; c < 8; c++) {
      data[dstOff + c] = block[srcOff + c];
    }
  }
}