export const MODES = Object.freeze({ broken: 128, fixed: 512, pooled: 8 });

// This lesson has a fixed Conv2d/ReLU/MaxPool2d stack, not a general shape engine.
export function inspectShapes({ batch, height, width, mode = 'broken' }) {
  const values = { batch, height, width };
  for (const [name, value] of Object.entries(values)) {
    const min = name === 'batch' ? 1 : 2;
    const max = name === 'batch' ? 64 : 256;
    if (!Number.isInteger(value) || value < min || value > max) {
      return { valid: false, message: `${name === 'batch' ? 'Batch size' : name === 'height' ? 'Height' : 'Width'} must be a whole number from ${min} to ${max}.` };
    }
  }
  if (!Object.hasOwn(MODES, mode)) return { valid: false, message: 'Choose one of the three classifier heads.' };

  const pooledHeight = Math.floor(height / 2);
  const pooledWidth = Math.floor(width / 2);
  const originalFeatures = 8 * pooledHeight * pooledWidth;
  const features = mode === 'pooled' ? 8 : originalFeatures;
  const inFeatures = MODES[mode];
  const matches = features === inFeatures;
  const reshapedBatch = batch * originalFeatures / 128;
  const canReshape = Number.isInteger(reshapedBatch);
  return {
    valid: true,
    batch, height, width, mode,
    input: [batch, 3, height, width],
    conv: [batch, 8, height, width],
    pool: [batch, 8, pooledHeight, pooledWidth],
    adaptive: mode === 'pooled' ? [batch, 8, 1, 1] : null,
    flattened: [batch, features],
    weight: [5, inFeatures],
    output: matches ? [batch, 5] : null,
    features, originalFeatures, inFeatures, matches,
    error: matches ? null : `mat1 and mat2 shapes cannot be multiplied (${batch}x${features} and ${inFeatures}x5)`,
    unsafeReshape: {
      possible: canReshape,
      shape: canReshape ? [reshapedBatch, 128] : null,
      output: canReshape ? [reshapedBatch, 5] : null,
      preservesBatch: canReshape && reshapedBatch === batch,
      elements: batch * originalFeatures,
    },
  };
}