// Dataset picker — ported from app/demo/index.html. Renders the six sklearn
// previews into cards, streams jittered random walks, and writes the current
// selection into hidden form inputs so the main <form> can submit it to the
// Prefect flow.
import * as THREE from 'three';
import { OrbitControls } from 'three/addons/controls/OrbitControls.js';

// Trajectories are precomputed at the max cycle length. Toggling n-frames
// truncates (12 is a prefix of 24 is a prefix of 48) so the same walk is
// reused — no reroll on toggle, and the per-frame pulse stays consistent.
const MAX_FRAMES = 48;

const CATEGORICAL_HEX = [
  '#1f4e5f', '#c97b3f', '#8b5a9f', '#5a8560',
  '#c74a5e', '#6b7d8f', '#b89f51', '#4a6fa5',
];
const CATEGORICAL = CATEGORICAL_HEX.map(h => new THREE.Color(h));

function rampContinuous(t) {
  const hue = (1 - t) * 215 + t * 28;
  const sat = 0.62;
  const lit = 0.50 + (t - 0.5) * 0.08;
  return new THREE.Color().setHSL(hue / 360, sat, lit);
}

function normalize(points) {
  const n = points.length;
  let mx = 0, my = 0, mz = 0;
  for (const p of points) { mx += p[0]; my += p[1]; mz += p[2]; }
  mx /= n; my /= n; mz /= n;

  // p95 of per-point max-coord magnitude — robust to distribution tails
  // (gaussian_quantiles / classification otherwise shrink to a fraction of
  // the viewport under pure max-abs normalization).
  const perPoint = new Float64Array(n);
  for (let i = 0; i < n; i++) {
    const p = points[i];
    const a = Math.abs(p[0] - mx);
    const b = Math.abs(p[1] - my);
    const c = Math.abs(p[2] - mz);
    perPoint[i] = a > b ? (a > c ? a : c) : (b > c ? b : c);
  }
  const sorted = Array.from(perPoint).sort((a, b) => a - b);
  const scale = Math.max(sorted[Math.floor(n * 0.95)], 1e-9);

  const out = new Float32Array(n * 3);
  for (let i = 0; i < n; i++) {
    out[i*3]   = (points[i][0] - mx) / scale;
    out[i*3+1] = (points[i][1] - my) / scale;
    out[i*3+2] = (points[i][2] - mz) / scale;
  }
  return { positions: out };
}

function buildColors(labels, kind) {
  const n = labels.length;
  const colors = new Float32Array(n * 3);
  if (kind === 'categorical') {
    for (let i = 0; i < n; i++) {
      const c = CATEGORICAL[labels[i] % CATEGORICAL.length];
      colors[i*3] = c.r; colors[i*3+1] = c.g; colors[i*3+2] = c.b;
    }
  } else {
    let lo = Infinity, hi = -Infinity;
    for (const v of labels) { if (v < lo) lo = v; if (v > hi) hi = v; }
    const range = (hi - lo) || 1;
    for (let i = 0; i < n; i++) {
      const c = rampContinuous((labels[i] - lo) / range);
      colors[i*3] = c.r; colors[i*3+1] = c.g; colors[i*3+2] = c.b;
    }
  }
  return colors;
}

function applyPickerBackground(scene, el) {
  const css = getComputedStyle(el).getPropertyValue('--picker-panel').trim();
  scene.background = new THREE.Color(css || '#f2eee4');
}

function createScene(container, dataset) {
  const { positions: basePositions } = normalize(dataset.points);
  const colors = buildColors(dataset.labels, dataset.kind);
  const positions = new Float32Array(basePositions);

  const geometry = new THREE.BufferGeometry();
  geometry.setAttribute('position', new THREE.BufferAttribute(positions, 3));
  geometry.setAttribute('color', new THREE.BufferAttribute(colors, 3));

  const material = new THREE.PointsMaterial({
    size: 2.1,
    sizeAttenuation: false,
    vertexColors: true,
    transparent: true,
    opacity: 0.92,
  });

  const scene = new THREE.Scene();
  applyPickerBackground(scene, container);
  scene.add(new THREE.Points(geometry, material));

  const camera = new THREE.PerspectiveCamera(42, 1, 0.1, 100);
  camera.position.set(2.6, 1.9, 2.6);
  camera.lookAt(0, 0, 0);

  const renderer = new THREE.WebGLRenderer({ antialias: true, alpha: false });
  renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
  container.appendChild(renderer.domElement);

  const controls = new OrbitControls(camera, renderer.domElement);
  controls.enableDamping = true;
  controls.dampingFactor = 0.08;
  controls.enablePan = false;
  controls.autoRotate = true;
  controls.autoRotateSpeed = 0.55;
  controls.minDistance = 1.5;
  controls.maxDistance = 6;

  return {
    scene, camera, renderer, controls, container, geometry,
    basePositions,
    trajectories: null,
    numFrames: 0,
    snapshotMs: 1000 / 12,
    // holdMs pads the end of each cycle at rest, so frame-0-as-rest is
    // actually visible (otherwise it's zero-duration).
    holdMs: 200,
    // Shared across scenes — applyF() resets all of them together so the
    // three previews stay in lockstep through n-frames toggles.
    cycleStartMs: 0,
  };
}

function buildTrajectories(s, numFrames) {
  const n = s.basePositions.length;
  const buf = new Float32Array(numFrames * n);
  for (let f = 1; f < numFrames; f++) {
    const prev = (f - 1) * n;
    const cur = f * n;
    for (let j = 0; j < n; j += 2) {
      const u1 = Math.random() || 1e-12;
      const u2 = Math.random();
      const mag = Math.sqrt(-2.0 * Math.log(u1));
      buf[cur + j] = buf[prev + j] + mag * Math.cos(2 * Math.PI * u2);
      if (j + 1 < n) buf[cur + j + 1] = buf[prev + j + 1] + mag * Math.sin(2 * Math.PI * u2);
    }
  }
  s.trajectories = buf;
  s.numFrames = numFrames;
}

function sizeScene(s) {
  const rect = s.container.getBoundingClientRect();
  const size = Math.max(1, Math.floor(rect.width));
  s.renderer.setSize(size, size);
  s.camera.aspect = 1;
  s.camera.updateProjectionMatrix();
}

async function main() {
  const gallery = document.getElementById('gallery');
  let data;
  try {
    const res = await fetch('/data.json');
    data = await res.json();
  } catch (err) {
    gallery.innerHTML =
      `<div class="picker-loading">failed to load /data.json — ${err}</div>`;
    return;
  }

  gallery.innerHTML = '';
  const scenes = [];
  const order = Object.entries(data);
  let selectedId = null;

  // Hidden form inputs the main <form> will submit.
  const hidden = {
    datasetId: document.getElementById('dataset_id'),
    numPoints: document.getElementById('num_points'),
    numTimesteps: document.getElementById('num_timesteps'),
    jitterScale: document.getElementById('jitter_scale'),
  };

  const pickerDetails = document.getElementById('picker');
  const summaryPath = document.getElementById('picker-summary-path');
  const selectedPath = document.getElementById('selected-path');
  const continueBtn = document.getElementById('continue-btn');

  const vizToScene = new WeakMap();
  const sizeObserver = new ResizeObserver((entries) => {
    for (const entry of entries) {
      const s = vizToScene.get(entry.target);
      if (s) sizeScene(s);
    }
  });

  order.forEach(([id, ds], i) => {
    const card = document.createElement('div');
    card.className = 'card';
    card.dataset.id = id;
    card.innerHTML = `
      <div class="viz">
        <span class="fig-label">Fig. 1.${i + 1}</span>
        <span class="key-hint">[${i + 1}]</span>
        <span class="controls-hint">drag · scroll</span>
      </div>
      <div class="card-body">
        <div class="card-label">
          <span class="dot"></span>
          <span>${ds.name}</span>
        </div>
        <div class="card-path">${ds.path}</div>
        <div class="card-desc">${ds.description}</div>
      </div>
    `;
    gallery.appendChild(card);

    const viz = card.querySelector('.viz');
    const s = createScene(viz, ds);
    buildTrajectories(s, MAX_FRAMES);
    scenes.push(s);
    vizToScene.set(viz, s);
    sizeObserver.observe(viz);

    s.controls.addEventListener('start', () => { s.controls.autoRotate = false; });

    card.addEventListener('click', () => selectCard(id, card, ds));
  });

  function selectCard(id, card, ds) {
    document.querySelectorAll('#gallery .card').forEach(c => c.classList.remove('selected'));
    card.classList.add('selected');
    selectedId = id;
    selectedPath.textContent = ds.path;
    hidden.datasetId.value = id;
    updateContinue();
  }

  document.addEventListener('themechange', () => {
    for (const s of scenes) applyPickerBackground(s.scene, s.container);
  });

  const formSubmitBtn = document.querySelector('#run-form button.submit');
  function updateContinue() {
    continueBtn.disabled = !selectedId;
    continueBtn.title = selectedId ? '' : 'pick a dataset first';
    if (formSubmitBtn) {
      formSubmitBtn.disabled = !selectedId;
      formSubmitBtn.title = selectedId ? '' : 'pick a dataset first';
    }
  }
  updateContinue();

  const nInputs = document.querySelectorAll('input[name="n"]');
  function applyN(n) {
    hidden.numPoints.value = String(n);
    for (const s of scenes) {
      const cap = s.geometry.attributes.position.count;
      s.geometry.setDrawRange(0, Math.min(n, cap));
    }
  }
  nInputs.forEach(input => {
    input.addEventListener('change', (e) => applyN(parseInt(e.target.value, 10)));
  });
  applyN(parseInt(document.querySelector('input[name="n"]:checked').value, 10));

  let jitterScale = 0;
  const jInputs = document.querySelectorAll('input[name="j"]');
  function applyJ(v) {
    jitterScale = v;
    hidden.jitterScale.value = String(v);
  }
  jInputs.forEach(input => {
    input.addEventListener('change', (e) => applyJ(parseFloat(e.target.value)));
  });
  applyJ(parseFloat(document.querySelector('input[name="j"]:checked').value));

  // timesteps: truncates the precomputed walk. Prefix-stable (12 ⊂ 24 ⊂ 48),
  // so toggling changes cycle length without rerolling. cycleStartMs is
  // shared so all cards animate in lockstep.
  const fInputs = document.querySelectorAll('input[name="f"]');
  function applyF(n) {
    hidden.numTimesteps.value = String(n);
    const start = performance.now();
    for (const s of scenes) {
      s.numFrames = n;
      s.cycleStartMs = start;
    }
  }
  fInputs.forEach(input => {
    input.addEventListener('change', (e) => applyF(parseInt(e.target.value, 10)));
  });
  applyF(parseInt(document.querySelector('input[name="f"]:checked').value, 10));

  function selectByIndex(idx, { scroll = true } = {}) {
    const entry = order[idx];
    if (!entry) return;
    const [id, ds] = entry;
    const card = gallery.children[idx];
    if (!card) return;
    selectCard(id, card, ds);
    if (scroll) card.scrollIntoView({ behavior: 'smooth', inline: 'nearest', block: 'nearest' });
  }
  document.addEventListener('keydown', (e) => {
    if (!pickerDetails.open) return;
    if (e.target.tagName === 'INPUT' || e.target.tagName === 'TEXTAREA') return;
    if (/^[1-9]$/.test(e.key)) {
      selectByIndex(parseInt(e.key, 10) - 1);
      return;
    }
    if (e.key === 'ArrowRight' || e.key === 'ArrowLeft') {
      e.preventDefault();
      const currentIdx = order.findIndex(([id]) => id === selectedId);
      const n = order.length;
      const nextIdx = e.key === 'ArrowRight'
        ? (currentIdx < 0 ? 0 : Math.min(currentIdx + 1, n - 1))
        : (currentIdx < 0 ? n - 1 : Math.max(currentIdx - 1, 0));
      selectByIndex(nextIdx);
    }
  });

  continueBtn.addEventListener('click', () => {
    if (!selectedId) return;
    const ds = data[selectedId];
    summaryPath.textContent = ds.path;
    pickerDetails.open = false;
  });

  function tick() {
    requestAnimationFrame(tick);
    // When the picker is collapsed the canvases are display:none inside a
    // closed <details>; rects are zero. Skip the per-frame work.
    if (!pickerDetails.open) return;
    const now = performance.now();
    for (const s of scenes) {
      const N = s.numFrames;
      const n = s.basePositions.length;
      const walkMs = N * s.snapshotMs;
      const cycleMs = walkMs + s.holdMs;
      const elapsed = ((now - s.cycleStartMs) % cycleMs + cycleMs) % cycleMs;

      const total = n / 3;
      const drawCount = s.geometry.drawRange.count;
      const visibleN = Number.isFinite(drawCount) ? Math.min(drawCount, total) : total;
      const limit = visibleN * 3;
      const pos = s.geometry.attributes.position.array;
      const base = s.basePositions;

      if (elapsed >= walkMs) {
        for (let i = 0; i < limit; i++) pos[i] = base[i];
      } else {
        const frameF = elapsed / s.snapshotMs;
        const frameIdx = Math.floor(frameF);
        const interpT = frameF - frameIdx;
        const nextIdx = (frameIdx + 1) % N;
        const aOff = frameIdx * n;
        const bOff = nextIdx * n;
        const tr = s.trajectories;
        const scale = jitterScale;
        const u = 1 - interpT;
        for (let i = 0; i < limit; i++) {
          pos[i] = base[i] + (tr[aOff + i] * u + tr[bOff + i] * interpT) * scale;
        }
      }
      s.geometry.attributes.position.needsUpdate = true;

      s.controls.update();
      s.renderer.render(s.scene, s.camera);
    }
  }
  tick();

  // Reopening the picker after it's been closed: canvases may have been
  // laid out at zero size while hidden. Re-measure on toggle.
  pickerDetails.addEventListener('toggle', () => {
    if (pickerDetails.open) {
      for (const s of scenes) sizeScene(s);
    }
  });
}

main();