const TOTAL_STEPS = 1000;
const modelInput = document.querySelector('#model-input');
const DEFAULT_CHECKPOINTS = [0, 50, 100, 200, 350, 500, 700, 1000];
const SIZE = 192;
const PIXEL_COUNT = SIZE * SIZE * 3;
const elements = {
  timestepsInput: document.querySelector('#timesteps-input'), applyTimesteps: document.querySelector('#apply-timesteps'),
  seedInput: document.querySelector('#seed-input'), randomizeSeed: document.querySelector('#randomize-seed'), configError: document.querySelector('#config-error'), diffusionTrack: document.querySelector('#diffusion-track'),
  denoisingTrack: document.querySelector('#denoising-track'), imageUpload: document.querySelector('#image-upload'), modelStatus: document.querySelector('#model-status'),
  forwardProgressLabel: document.querySelector('#forward-progress-label'), reverseProgressLabel: document.querySelector('#reverse-progress-label'),
};
let checkpoints = [...DEFAULT_CHECKPOINTS];
let sourcePixels = new Float32Array(PIXEL_COUNT);
let forwardStates = [];
let reverseStates = new Map();
let imageLabel = 'Still life 01';
let forwardRevealCount = 1;
let forwardCompleted = false;
let reverseNoiseState = null;
let reverseInputMode = 'seed';
let animationTimer = null;
let generationController = null;
let generationBusy = false;
const progressStartedAt = { forward: null, reverse: null };
const schedule = makeSchedule();

function makeSchedule() {
  const alphaBars = new Float64Array(TOTAL_STEPS); let product = 1;
  for (let step = 0; step < TOTAL_STEPS; step += 1) { const beta = .0001 + (.02 - .0001) * step / (TOTAL_STEPS - 1); product *= 1 - beta; alphaBars[step] = product; }
  return alphaBars;
}
function randomNormal(seed) {
  let state = seed >>> 0;
  return () => { state = (state + 0x6D2B79F5) | 0; let value = state; value = Math.imul(value ^ value >>> 15, value | 1); value ^= value + Math.imul(value ^ value >>> 7, value | 61); const uniform = ((value ^ value >>> 14) >>> 0) / 4294967296; return Math.sqrt(-2 * Math.log(Math.max(uniform, 1e-7))) * Math.cos(2 * Math.PI * uniform); };
}
function makeSourceArt() {
  const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; const context = canvas.getContext('2d');
  const sky = context.createLinearGradient(0, 0, SIZE, SIZE); sky.addColorStop(0, '#dce7db'); sky.addColorStop(.58, '#b8c7ad'); sky.addColorStop(1, '#a2b5a0'); context.fillStyle = sky; context.fillRect(0, 0, SIZE, SIZE);
  context.fillStyle = '#e9d6b5'; context.beginPath(); context.arc(147, 44, 24, 0, Math.PI * 2); context.fill();
  context.fillStyle = '#52765f'; context.beginPath(); context.moveTo(0, 137); context.quadraticCurveTo(50, 91, 99, 137); context.quadraticCurveTo(145, 109, 192, 132); context.lineTo(192, 192); context.lineTo(0, 192); context.fill();
  context.fillStyle = '#315b48'; context.beginPath(); context.moveTo(0, 159); context.quadraticCurveTo(61, 127, 112, 164); context.quadraticCurveTo(155, 143, 192, 157); context.lineTo(192, 192); context.lineTo(0, 192); context.fill();
  context.fillStyle = '#d8795d'; context.beginPath(); context.ellipse(92, 144, 24, 31, -.22, 0, Math.PI * 2); context.fill(); context.fillStyle = '#e8a37c'; context.beginPath(); context.ellipse(84, 136, 14, 21, -.22, 0, Math.PI * 2); context.fill();
  context.fillStyle = '#496c52'; context.beginPath(); context.ellipse(103, 111, 21, 10, -.62, 0, Math.PI * 2); context.fill(); context.beginPath(); context.ellipse(77, 109, 16, 8, .52, 0, Math.PI * 2); context.fill(); return canvas;
}
function canvasPixels(canvas) {
  const rgba = canvas.getContext('2d', { willReadFrequently: true }).getImageData(0, 0, SIZE, SIZE).data; const pixels = new Float32Array(PIXEL_COUNT);
  for (let pixel = 0; pixel < SIZE * SIZE; pixel += 1) for (let channel = 0; channel < 3; channel += 1) pixels[pixel * 3 + channel] = rgba[pixel * 4 + channel] / 127.5 - 1;
  return pixels;
}
function pixelsToCanvas(pixels, canvas) {
  const context = canvas.getContext('2d'); const image = context.createImageData(SIZE, SIZE);
  for (let index = 0; index < PIXEL_COUNT; index += 1) image.data[Math.floor(index / 3) * 4 + index % 3] = Math.round((Math.max(-1, Math.min(1, pixels[index])) + 1) * 127.5);
  for (let index = 3; index < image.data.length; index += 4) image.data[index] = 255; context.putImageData(image, 0, 0);
}
function drawState(canvas, state) { const context = canvas.getContext('2d'); context.clearRect(0, 0, SIZE, SIZE); if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0, SIZE, SIZE); else if (state instanceof Float32Array) pixelsToCanvas(state, canvas); }
function computeForwardStates() {
  const random = randomNormal(readSeed()); const noise = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) noise[index] = random();
  forwardStates = checkpoints.map((step) => { if (step === 0) return new Float32Array(sourcePixels); const alphaBar = schedule[step - 1]; const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = Math.sqrt(alphaBar) * sourcePixels[index] + Math.sqrt(1 - alphaBar) * noise[index]; return state; });
}
function makeNoiseState(seed) { const random = randomNormal(seed); const state = new Float32Array(PIXEL_COUNT); for (let index = 0; index < PIXEL_COUNT; index += 1) state[index] = random(); return state; }
function stateToPng(state) { const canvas = document.createElement('canvas'); const size = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256; canvas.width = size; canvas.height = size; const small = document.createElement('canvas'); small.width = SIZE; small.height = SIZE; drawState(small, state); canvas.getContext('2d').drawImage(small, 0, 0, size, size); return canvas.toDataURL('image/png'); }
function setModelStatus(message) { if (elements.modelStatus) elements.modelStatus.textContent = message; }
function resampleNoise() {
  const values = new Uint32Array(1); crypto.getRandomValues(values); const seed = values[0] & 0x7fffffff; elements.seedInput.value = String(seed); reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(seed); reverseStates = new Map(); resetProgress('reverse'); setModelStatus(`Resampled Gaussian noise · seed ${seed}.`); renderTracks();
}
function useForwardNoise() {
  if (!forwardCompleted) return; reverseInputMode = 'forward'; reverseNoiseState = new Float32Array(forwardStates.at(-1)); reverseStates = new Map(); resetProgress('reverse'); setModelStatus('Using the final forward-process noise as the reverse input.'); renderTracks();
}
function makeActionButton(label, id, action) { const button = document.createElement('button'); button.type = 'button'; button.id = id; button.className = 'process-button'; button.textContent = label; button.addEventListener('click', action); return button; }
function makeCard(step, state, type, index) {
  const card = document.createElement('article'); card.className = 'stage-card'; card.dataset.step = String(step); const actionSlot = document.createElement('div'); actionSlot.className = 'action-slot';
  if (type === 'forward' && index >= forwardRevealCount) card.classList.add('is-pending');
  if (type === 'forward' && index === 0) actionSlot.append(makeActionButton('Diffuse', 'diffuse-button', runDiffusion));
  if (type === 'forward' && index === checkpoints.length - 1) { actionSlot.classList.add('status-right'); actionSlot.append(elements.forwardProgressLabel); }
  if (type === 'reverse' && index === 0) actionSlot.append(elements.reverseProgressLabel);
  if (type === 'reverse' && index === checkpoints.length - 1) {
    actionSlot.classList.add('reverse-actions');
    actionSlot.append(makeActionButton('Denoise', 'denoise-button', generateReverse));
  }
  const canvasWrap = document.createElement('div'); canvasWrap.className = 'stage-canvas'; const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; drawState(canvas, state); canvasWrap.append(canvas);
  if (type === 'forward' && index === 0) { const change = document.createElement('button'); change.className = 'change-image'; change.type = 'button'; change.textContent = 'Change image'; change.addEventListener('click', () => elements.imageUpload.click()); canvasWrap.append(change); }
  if (type === 'reverse' && index === checkpoints.length - 1) { const noiseActions = document.createElement('div'); noiseActions.className = 'noise-actions'; noiseActions.append(makeActionButton('Resample', 'resample-button', resampleNoise), makeActionButton('Forward noise', 'forward-noise-button', useForwardNoise)); canvasWrap.append(noiseActions); }
  const caption = document.createElement('div'); caption.className = 'stage-caption';
  const titleText = type === 'forward' && index === 0 ? imageLabel : step === 0 ? 'Generated image' : step === TOTAL_STEPS ? 'Gaussian noise' : '';
  if (titleText) { const title = document.createElement('strong'); title.textContent = titleText; caption.append(title); }
  const time = document.createElement('span'); time.textContent = `t = ${step}`; caption.append(time);
  const preview = document.createElement('div'); preview.className = 'hover-preview';
  const previewCanvas = document.createElement('canvas');
  const nativeSize = modelInput.value === 'google/ddpm-cifar10-32' ? 32 : 256;
  previewCanvas.width = state instanceof HTMLCanvasElement ? state.width : nativeSize;
  previewCanvas.height = state instanceof HTMLCanvasElement ? state.height : nativeSize;
  preview.style.width = `${previewCanvas.width + 4}px`; preview.style.height = `${previewCanvas.height + 4}px`;
  previewCanvas.style.width = `${previewCanvas.width}px`; previewCanvas.style.height = `${previewCanvas.height}px`;
  if (state) {
    const context = previewCanvas.getContext('2d');
    if (state instanceof HTMLCanvasElement) context.drawImage(state, 0, 0);
    else { const temporary = document.createElement('canvas'); temporary.width = SIZE; temporary.height = SIZE; drawState(temporary, state); context.drawImage(temporary, 0, 0, previewCanvas.width, previewCanvas.height); }
  }
  preview.append(previewCanvas);
  canvas.addEventListener('click', () => {
    document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
    const rect = canvasWrap.getBoundingClientRect();
    preview.style.left = `${Math.max(8, Math.min(window.innerWidth - previewCanvas.width - 12, rect.left + rect.width / 2 - previewCanvas.width / 2))}px`;
    preview.style.top = `${Math.max(8, Math.min(window.innerHeight - previewCanvas.height - 12, rect.top + rect.height / 2 - previewCanvas.height / 2))}px`;
    preview.classList.add('is-visible');
  });
  card.append(actionSlot, canvasWrap, caption, preview); return card;
}
function renderTracks() {
  elements.diffusionTrack.style.setProperty('--stage-count', checkpoints.length);
  elements.denoisingTrack.style.setProperty('--stage-count', checkpoints.length);
  elements.diffusionTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, forwardStates[index], 'forward', index)));
  elements.denoisingTrack.replaceChildren(...checkpoints.map((step, index) => makeCard(step, reverseStates.get(step) || (step === TOTAL_STEPS ? reverseNoiseState : null), 'reverse', index)));
  modelInput.disabled = generationBusy;
  const button = document.querySelector('#denoise-button'); if (button) { button.disabled = generationBusy; button.textContent = generationBusy ? 'Denoising…' : 'Denoise'; }
  const forwardNoiseButton = document.querySelector('#forward-noise-button'); if (forwardNoiseButton) { forwardNoiseButton.disabled = generationBusy || !forwardCompleted; forwardNoiseButton.title = forwardCompleted ? 'Use the last forward-process image' : 'Run Diffuse first'; }
  const resampleButton = document.querySelector('#resample-button'); if (resampleButton) resampleButton.disabled = generationBusy;
}
function formatDuration(seconds) {
  if (!Number.isFinite(seconds)) return '?';
  const whole = Math.max(0, Math.round(seconds)); const hours = Math.floor(whole / 3600); const minutes = Math.floor(whole % 3600 / 60); const remainder = whole % 60;
  return hours ? `${String(hours).padStart(2, '0')}:${String(minutes).padStart(2, '0')}:${String(remainder).padStart(2, '0')}` : `${String(minutes).padStart(2, '0')}:${String(remainder).padStart(2, '0')}`;
}
function resetProgress(kind) { progressStartedAt[kind] = null; setProgress(kind, 0); }
function beginProgress(kind) { progressStartedAt[kind] = performance.now(); setProgress(kind, 0); }
function setProgress(kind, completedSteps) {
  const completed = Math.max(0, Math.min(TOTAL_STEPS, Math.round(completedSteps))); const label = kind === 'forward' ? elements.forwardProgressLabel : elements.reverseProgressLabel;
  const elapsed = progressStartedAt[kind] === null ? 0 : (performance.now() - progressStartedAt[kind]) / 1000; const rate = elapsed > 0 && completed > 0 ? completed / elapsed : 0; const remaining = rate > 0 ? (TOTAL_STEPS - completed) / rate : Infinity;
  const text = `${completed}/${TOTAL_STEPS} [${formatDuration(elapsed)}<${completed === TOTAL_STEPS ? '00:00' : formatDuration(remaining)}, ${rate > 0 ? rate.toFixed(2) : '?'}it/s]`; label.value = text; label.textContent = text;
}
function stopAnimation() { if (animationTimer) window.clearInterval(animationTimer); animationTimer = null; document.querySelectorAll('.stage-card.is-active').forEach((card) => card.classList.remove('is-active')); }
function runDiffusion() {
  try { computeForwardStates(); elements.configError.textContent = ''; } catch (error) { elements.configError.textContent = error.message; return; }
  stopAnimation(); forwardCompleted = false; forwardRevealCount = 1; beginProgress('forward'); renderTracks(); const cards = [...elements.diffusionTrack.querySelectorAll('.stage-card')]; cards[0].classList.add('is-active');
  const duration = Math.max(1000, (checkpoints.length - 1) * 650); const startedAt = performance.now(); let activeIndex = 0;
  animationTimer = window.setInterval(() => {
    const completed = Math.min(TOTAL_STEPS, Math.floor((performance.now() - startedAt) / duration * TOTAL_STEPS)); setProgress('forward', completed);
    let nextIndex = activeIndex; while (nextIndex + 1 < checkpoints.length && checkpoints[nextIndex + 1] <= completed) nextIndex += 1;
    if (nextIndex !== activeIndex) { cards[activeIndex].classList.remove('is-active'); for (let index = activeIndex + 1; index <= nextIndex; index += 1) cards[index].classList.remove('is-pending'); cards[nextIndex].classList.add('is-active'); activeIndex = nextIndex; forwardRevealCount = nextIndex + 1; }
    if (completed >= TOTAL_STEPS) { forwardCompleted = true; setProgress('forward', TOTAL_STEPS); stopAnimation(); renderTracks(); }
  }, 50);
}
function parseTimesteps() {
  const values = elements.timestepsInput.value.split(',').map((part) => part.trim()).filter(Boolean).map(Number);
  if (values.length < 2 || values.some((value) => !Number.isInteger(value) || value < 0 || value > TOTAL_STEPS)) throw new Error('Enter at least two whole numbers between 0 and 1000.');
  const unique = [...new Set(values)].sort((a, b) => a - b); if (unique[0] !== 0 || unique.at(-1) !== TOTAL_STEPS) throw new Error('Include both 0 and 1000 so each process has a start and end.'); if (unique.length > 12) throw new Error('Use no more than 12 timesteps.'); return unique;
}
function readSeed() {
  const value = Number(elements.seedInput.value.trim());
  if (!Number.isSafeInteger(value) || value < 0 || value > 2147483647) throw new Error('Seed must be a whole number from 0 to 2147483647.');
  return value;
}
function applyTimesteps() {
  try { const next = parseTimesteps(); stopAnimation(); if (generationController) generationController.abort(); checkpoints = next; reverseStates = new Map(); reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(readSeed()); forwardRevealCount = 1; forwardCompleted = false; computeForwardStates(); resetProgress('forward'); resetProgress('reverse'); elements.timestepsInput.value = checkpoints.join(', '); elements.configError.textContent = ''; setModelStatus('Timesteps updated. Select Denoise to run the model.'); renderTracks(); }
  catch (error) { elements.configError.textContent = error.message; }
}
function decodeFrame(data) { return new Promise((resolve, reject) => { const image = new Image(); image.onload = () => { const canvas = document.createElement('canvas'); canvas.width = image.naturalWidth; canvas.height = image.naturalHeight; canvas.getContext('2d').drawImage(image, 0, 0); resolve(canvas); }; image.onerror = () => reject(new Error('The backend returned an unreadable image frame.')); image.src = `data:image/jpeg;base64,${data}`; }); }
async function generateReverse() {
  if (generationBusy) return;
  let seed;
  try { seed = readSeed(); elements.configError.textContent = ''; } catch (error) { elements.configError.textContent = error.message; return; }
  stopAnimation(); generationController = new AbortController(); generationBusy = true; reverseStates = new Map(); resetProgress('reverse'); setModelStatus('Loading model weights if needed, then sampling 1,000 reverse steps…'); renderTracks();
  try {
    const initialImage = reverseInputMode === 'forward' ? stateToPng(reverseNoiseState) : null;
    const response = await fetch('/api/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ model: modelInput.value, checkpoints, seed, initial_image: initialImage }), signal: generationController.signal }); if (!response.ok || !response.body) throw new Error(`Backend request failed (${response.status}).`);
    const reader = response.body.getReader(); const decoder = new TextDecoder(); let pending = '';
    while (true) { const { value, done } = await reader.read(); pending += decoder.decode(value || new Uint8Array(), { stream: !done }); const lines = pending.split('\n'); pending = lines.pop(); for (const line of lines) if (line.trim()) await handleModelEvent(JSON.parse(line)); if (done) break; }
    if (pending.trim()) await handleModelEvent(JSON.parse(pending));
  } catch (error) { if (error.name !== 'AbortError') setModelStatus(`Model error: ${error.message}`); }
  finally { generationBusy = false; generationController = null; renderTracks(); }
}
async function handleModelEvent(event) {
  if (event.type === 'status') setModelStatus(`Loading ${event.model} on the first generation…`);
  if (event.type === 'progress') { const completed = event.total - event.step; if (progressStartedAt.reverse === null) beginProgress('reverse'); setProgress('reverse', completed); setModelStatus(`Denoising · ${completed}/${event.total} steps complete.`); }
  if (event.type === 'frame') { reverseStates.set(event.step, await decodeFrame(event.image)); setProgress('reverse', TOTAL_STEPS - event.step); setModelStatus(`Real model inference · received t = ${event.step}.`); renderTracks(); const card = elements.denoisingTrack.querySelector(`[data-step="${event.step}"]`); if (card) { document.querySelectorAll('#denoising-track .is-active').forEach((item) => item.classList.remove('is-active')); card.classList.add('is-active'); } }
  if (event.type === 'complete') { setProgress('reverse', TOTAL_STEPS); setModelStatus(`Generated with ${event.model || modelInput.value} on ${event.device.toUpperCase()} · seed ${event.seed}.`); } if (event.type === 'error') throw new Error(event.message);
}
function loadImage(image, label) { const canvas = document.createElement('canvas'); canvas.width = SIZE; canvas.height = SIZE; const context = canvas.getContext('2d'); const scale = Math.max(SIZE / image.width, SIZE / image.height); const width = image.width * scale; const height = image.height * scale; context.drawImage(image, (SIZE - width) / 2, (SIZE - height) / 2, width, height); sourcePixels = canvasPixels(canvas); imageLabel = label; reverseInputMode = 'seed'; reverseNoiseState = makeNoiseState(readSeed()); forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); }
modelInput.addEventListener('change', () => {
  applyTimesteps();
  const selectedModel = modelInput.value;
  const source = selectedModel === 'google/ddpm-cifar10-32'
    ? { path: 'assets/cifar10_train_23_32.png', label: 'CIFAR-10 train 23' }
    : { path: 'assets/celeba_hq_validation_91_256.png', label: 'Human face' };
  const image = new Image();
  image.onload = () => {
    if (modelInput.value === selectedModel) loadImage(image, source.label);
  };
  image.onerror = () => setModelStatus('Could not load the selected model’s source image.');
  image.src = source.path;
});
elements.applyTimesteps.addEventListener('click', applyTimesteps); elements.timestepsInput.addEventListener('keydown', (event) => { if (event.key === 'Enter') applyTimesteps(); });
elements.randomizeSeed.addEventListener('click', () => { resampleNoise(); elements.configError.textContent = ''; forwardRevealCount = 1; forwardCompleted = false; resetProgress('forward'); computeForwardStates(); renderTracks(); });
elements.imageUpload.addEventListener('change', (event) => { const file = event.target.files?.[0]; if (!file) return; const image = new Image(); image.onload = () => { loadImage(image, file.name.replace(/\.[^.]+$/, '')); URL.revokeObjectURL(image.src); }; image.src = URL.createObjectURL(file); });
// Keep the original generated still life available while the default face loads,
// and as a fallback if the image asset cannot be read.
sourcePixels = canvasPixels(makeSourceArt()); computeForwardStates(); reverseNoiseState = makeNoiseState(readSeed()); resetProgress('forward'); resetProgress('reverse'); renderTracks();
const defaultSourceImage = new Image();
defaultSourceImage.onload = () => {
  if (modelInput.value === 'google/ddpm-celebahq-256') loadImage(defaultSourceImage, 'Human face');
};
defaultSourceImage.src = 'assets/celeba_hq_validation_91_256.png';

document.addEventListener('pointerdown', event => {
  if (!event.target.closest('.stage-canvas canvas')) document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
});
document.addEventListener('keydown', event => {
  if (event.key === 'Escape') document.querySelectorAll('.hover-preview.is-visible').forEach(item => item.classList.remove('is-visible'));
});