import { inspectShapes } from './shape-math.mjs';

const byId = (id) => document.getElementById(id);
const shape = (dimensions) => `[${dimensions.join(', ')}]`;
const text = (id, value) => { byId(id).textContent = value; };
const form = byId('shape-form');

function render() {
  const result = inspectShapes({
    batch: byId('batch').valueAsNumber,
    height: byId('height').valueAsNumber,
    width: byId('width').valueAsNumber,
    mode: form.elements.mode.value,
  });
  byId('invalid').hidden = result.valid;
  byId('calculation').hidden = !result.valid;
  byId('reshape-result').hidden = !result.valid;
  if (!result.valid) { text('invalid', result.message); return; }

  const r = result;
  text('shape-input', shape(r.input));
  text('shape-conv', shape(r.conv));
  text('shape-pool', shape(r.pool));
  text('shape-flat', shape(r.flattened));
  byId('adaptive-row').hidden = !r.adaptive;
  if (r.adaptive) text('shape-adaptive', shape(r.adaptive));
  const linearName = byId('linear-name');
  linearName.replaceChildren(document.createTextNode(`Linear(${r.inFeatures}, 5)`));
  const small = document.createElement('small');
  small.textContent = `Expects ${r.inFeatures} features per image`;
  linearName.append(small);
  text('shape-output', r.output ? shape(r.output) : `${r.features} ≠ ${r.inFeatures}`);
  byId('linear-row').className = r.matches ? 'match' : 'mismatch';
  byId('verdict').className = `verdict ${r.matches ? 'match' : 'mismatch'}`;
  text('verdict-title', r.matches ? `The widths match. Batch stays ${r.batch}.` : `${r.features} features arrived. The head expects ${r.inFeatures}.`);
  text('feature-math', r.adaptive
    ? '8 channels × 1 height × 1 width = 8 features per image after adaptive pooling.'
    : `8 channels × ${r.pool[2]} height × ${r.pool[3]} width = ${r.features} features per image.`);
  text('matrix-left', shape(r.flattened));
  text('matrix-right', shape([r.inFeatures, 5]));
  text('matrix-answer', r.output ? `= ${shape(r.output)}` : 'incompatible');
  text('weight-shape', shape(r.weight));
  byId('error-message').hidden = r.matches;
  text('error-message', r.error ?? '');
  const notes = {
    broken: r.matches
      ? 'This input produces exactly 128 features per image. Try 16 × 16: the same fixed head will no longer fit.'
      : 'The fixed head expects 128 features per image. Match the head to the intended input contract; changing the batch does not change this width.',
    fixed: r.matches
      ? '512 fits the 16 × 16 example. Try a wider image: changing the input can break a fixed-width head again.'
      : '512 is fixed for the original 16 × 16 example. This input has a different feature width. A successful earlier capture does not cover every input.',
    pooled: 'Adaptive pooling averages each spatial channel to one value, so this head receives 8 features. This changes the architecture and the classifier weights; it is not the same computation as the fixed-width model.',
  };
  text('mode-note', notes[r.mode]);
  text('reshape-before', shape([r.batch, r.originalFeatures]));
  const unsafe = r.unsafeReshape;
  text('reshape-after', unsafe.possible ? shape(unsafe.shape) : 'cannot reshape');
  const warning = byId('reshape-warning');
  warning.replaceChildren();
  const strong = document.createElement('strong');
  strong.textContent = !unsafe.possible
    ? `${unsafe.elements} elements cannot fill rows of 128.`
    : unsafe.preservesBatch
      ? `The feature width was already 128.`
      : `${r.batch} ${r.batch === 1 ? 'image became' : 'images became'} ${unsafe.shape[0]} ${unsafe.shape[0] === 1 ? 'row' : 'rows'}.`;
  warning.append(strong, document.createTextNode(!unsafe.possible
    ? ' This reshape would fail before Linear. Use a batch-preserving flatten and inspect the feature width.'
    : unsafe.preservesBatch
      ? ' This reshape is redundant here. It does not repair a mismatch or make other image sizes safe.'
      : ` Linear can now return ${shape(unsafe.output)}, but that is not one prediction per image.`));
}

form.addEventListener('input', render);
form.addEventListener('change', render);
form.addEventListener('submit', (event) => event.preventDefault());
document.querySelectorAll('[data-height]').forEach((button) => button.addEventListener('click', () => {
  byId('height').value = button.dataset.height;
  byId('width').value = button.dataset.width;
  render();
}));
byId('reset').addEventListener('click', () => { form.reset(); render(); });
render();