1import { monaco, type CustomEditorProvider, type Workbench, type WorkspaceFileSystem } from 'minwebide';
2import { getRunState, onDidChangeRunState, type RunState } from './runEvents';
3import { dirnameOf, outputDirFor, parseSampleFile, samplingDefaults, updateSampleYaml } from './sampleConfig';
4import './sampleEditor.css';
6// The default view for .sample files: a form over the YAML (shared text
7// model, so 'Reopen as Text Editor', dirty state, and Ctrl+S behave), plus
8// the run button and per-chain progress bars fed by runEvents.
10interface StopHandle {
11 stop(): void;
12}
14export function createSampleEditorProvider(fs: WorkspaceFileSystem, workbench: Workbench, stopHandle: StopHandle): CustomEditorProvider {
15 return {
16 viewType: 'stan.sampleView',
17 displayName: 'Sampling Run',
18 selector: [{ filenamePattern: '*.sample' }],
19 priority: 'default',
20 async resolveCustomEditor(doc) {
21 const model = await doc.getTextModel();
22 const uriKey = doc.uri.toString();
23 const sampleDir = dirnameOf(doc.uri.path);
24 const disposables: { dispose(): void }[] = [];
26 const element = el('div', 'sample-editor');
27 const inner = el('div', 'sample-editor-inner');
28 element.appendChild(inner);
30 const fileName = doc.uri.path.split('/').pop() ?? doc.uri.path;
31 inner.appendChild(el('h2', undefined, fileName));
32 inner.appendChild(el('p', 'sample-editor-subtitle', 'A sampling run: the Stan program, the data, and sampling parameters. This form edits the underlying YAML (tab menu → Reopen as Text Editor).'));
34 const problems = el('div', 'sample-problems');
35 problems.style.display = 'none';
36 inner.appendChild(problems);
38 // --- run button + progress (at the top, above the fields) ---------
39 const runRow = el('div', 'sample-run-row');
40 const runButton = el('button', 'sample-run-button', 'Run sampling');
41 runButton.addEventListener('click', () => {
42 const state = getRunState(uriKey);
43 if (isRunning(state)) {
44 stopHandle.stop();
45 } else {
46 void workbench.runFile(doc.uri);
47 }
48 });
49 // shown when the derived output dir holds a completed run
50 const runJsonPath = `${outputDirFor(doc.uri.path)}/run.json`;
51 const resultsButton = el('button', 'sample-results-button', 'View results');
52 resultsButton.style.display = 'none';
53 resultsButton.addEventListener('click', () => {
54 void workbench.openFile(fs.root.with({ path: runJsonPath }));
55 });
56 const runStatus = el('span', 'sample-run-status', '');
57 runRow.append(runButton, resultsButton, runStatus);
58 inner.appendChild(runRow);
60 inner.appendChild(el('p', 'sample-output-note',
61 `results are written to ${referenceFor(outputDirFor(doc.uri.path))} (replaced on each run)`));
63 const chainsBox = el('div', 'sample-chains');
64 inner.appendChild(chainsBox);
66 // --- fields ------------------------------------------------------
67 let applyingEdit = false;
68 const setKey = (key: string, value: string | number | undefined) => {
69 const updated = updateSampleYaml(model.getValue(), key, value);
70 if (updated !== model.getValue()) {
71 applyingEdit = true;
72 try {
73 model.pushEditOperations([], [{ range: model.getFullModelRange(), text: updated }], () => null);
74 } finally {
75 applyingEdit = false;
76 }
77 refresh();
78 }
79 };
81 const stanField = fileSelect('stan', 'the Stan program', '.stan');
82 const dataField = fileSelect('data', 'the data (JSON)', '.json');
84 const params = el('div', 'sample-params');
85 inner.appendChild(params);
86 const numberField = (key: 'num_chains' | 'num_warmup' | 'num_samples' | 'init_radius' | 'seed', hint: string, opts: { min: number; max?: number; step?: string; optional?: boolean }) => {
87 const input = el('input');
88 input.type = 'number';
89 input.min = String(opts.min);
90 if (opts.max !== undefined) {
91 input.max = String(opts.max);
92 }
93 input.step = opts.step ?? '1';
94 if (opts.optional) {
95 input.placeholder = 'random';
96 }
97 input.addEventListener('change', () => {
98 const raw = input.value.trim();
99 if (!raw) {
100 setKey(key, opts.optional ? undefined : samplingDefaults[key as keyof typeof samplingDefaults]);
101 return;
102 }
103 const value = Number(raw);
104 if (Number.isFinite(value)) {
105 setKey(key, value);
106 }
107 });
108 params.appendChild(field(key, hint, input, true));
109 return input;
110 };
111 const chainsInput = numberField('num_chains', 'chains', { min: 1, max: 8 });
112 const warmupInput = numberField('num_warmup', 'warmup iterations', { min: 0 });
113 const samplesInput = numberField('num_samples', 'draws per chain', { min: 1 });
114 const radiusInput = numberField('init_radius', 'init radius', { min: 0, step: '0.1' });
115 const seedInput = numberField('seed', 'random seed', { min: 0, optional: true });
117 const renderRunState = (state: RunState) => {
118 const running = isRunning(state);
119 runButton.textContent = running ? 'Stop' : 'Run sampling';
120 runButton.classList.toggle('stop', running);
121 runStatus.textContent = state.message ?? '';
122 runStatus.className = 'sample-run-status'
123 + (state.phase === 'failed' ? ' error' : state.phase === 'done' ? ' done' : '');
124 chainsBox.textContent = '';
125 if (state.chains) {
126 state.chains.forEach((chain, index) => {
127 const row = el('div', 'sample-chain');
128 const percent = chain.totalIterations > 0 ? Math.round((chain.iteration / chain.totalIterations) * 100) : 0;
129 row.appendChild(el('span', 'sample-chain-label',
130 `Chain ${index + 1} ${chain.iteration} / ${chain.totalIterations}${chain.iteration > 0 ? (chain.warmup ? ' (warmup)' : ' (sampling)') : ''}`));
131 const bar = el('div', 'sample-chain-bar');
132 const fill = el('div', 'sample-chain-fill' + (chain.warmup ? ' warmup' : ''));
133 fill.style.width = `${percent}%`;
134 bar.appendChild(fill);
135 row.appendChild(bar);
136 chainsBox.appendChild(row);
137 });
138 }
139 };
140 renderRunState(getRunState(uriKey));
141 disposables.push(onDidChangeRunState((key, state) => {
142 if (key === uriKey) {
143 renderRunState(state);
144 }
145 }));
147 // --- model → form ------------------------------------------------
148 const refresh = () => {
149 const { config, errors, warnings } = parseSampleFile(model.getValue());
150 const messages = [...errors, ...warnings.map(w => `warning: ${w}`)];
151 problems.style.display = messages.length ? '' : 'none';
152 problems.className = 'sample-problems' + (errors.length ? '' : ' warnings');
153 problems.textContent = messages.join('\n');
154 runButton.disabled = errors.length > 0;
156 setIfNotFocused(stanField.select, config.stan ?? '');
157 setIfNotFocused(dataField.select, config.data ?? '');
158 setIfNotFocused(chainsInput, String(config.num_chains));
159 setIfNotFocused(warmupInput, String(config.num_warmup));
160 setIfNotFocused(samplesInput, String(config.num_samples));
161 setIfNotFocused(radiusInput, String(config.init_radius));
162 setIfNotFocused(seedInput, config.seed === undefined ? '' : String(config.seed));
163 };
165 disposables.push(model.onDidChangeContent(() => {
166 if (!applyingEdit) {
167 refresh();
168 }
169 }));
171 // keep the .stan/.json dropdowns in sync with the project's files
172 const refreshFileLists = async () => {
173 const all = await listProjectFiles(fs);
174 stanField.setOptions(all.filter(path => path.endsWith('.stan')));
175 dataField.setOptions(all.filter(path => path.endsWith('.json')));
176 resultsButton.style.display = all.includes(runJsonPath) ? '' : 'none';
177 refresh();
178 };
179 let fileListTimer: ReturnType<typeof setTimeout> | undefined;
180 disposables.push(fs.fileService.onDidFilesChange(() => {
181 clearTimeout(fileListTimer);
182 fileListTimer = setTimeout(() => void refreshFileLists(), 300);
183 }));
184 await refreshFileLists();
186 return {
187 element,
188 dispose() {
189 clearTimeout(fileListTimer);
190 for (const disposable of disposables) {
191 disposable.dispose();
192 }
193 },
194 };
196 // --- helpers scoped to this pane ---------------------------------
198 /** A <select> of project files with a given extension, storing
199 * .sample-dir-relative references in the YAML. */
200 function fileSelect(key: 'stan' | 'data', hint: string, extension: string) {
201 const select = el('select');
202 let options: string[] = [];
203 const setOptions = (paths: string[]) => {
204 options = paths.map(path => referenceFor(path));
205 renderOptions(select.value);
206 };
207 const renderOptions = (current: string) => {
208 select.textContent = '';
209 const empty = makeOption('', `— select a ${extension} file —`);
210 select.appendChild(empty);
211 const seen = new Set<string>();
212 for (const reference of options) {
213 seen.add(reference);
214 select.appendChild(makeOption(reference, reference));
215 }
216 if (current && !seen.has(current)) {
217 select.appendChild(makeOption(current, `${current} (missing)`));
218 }
219 select.value = current;
220 };
221 select.addEventListener('change', () => setKey(key, select.value || undefined));
222 inner.appendChild(field(key, hint, select));
223 return { select, setOptions };
224 }
226 function referenceFor(path: string): string {
227 return path.startsWith(`${sampleDir}/`) && sampleDir !== '/'
228 ? path.slice(sampleDir.length + 1)
229 : (sampleDir === '/' ? path.slice(1) : path);
230 }
232 function setIfNotFocused(input: HTMLInputElement | HTMLSelectElement, value: string): void {
233 if (document.activeElement === input) {
234 return;
235 }
236 if (input instanceof HTMLSelectElement) {
237 const has = [...input.options].some(option => option.value === value);
238 if (!has && value) {
239 input.appendChild(makeOption(value, `${value} (missing)`));
240 }
241 }
242 if (input.value !== value) {
243 input.value = value;
244 }
245 }
247 function makeOption(value: string, label: string): HTMLOptionElement {
248 const option = document.createElement('option');
249 option.value = value;
250 option.textContent = label;
251 return option;
252 }
253 },
254 };
256 function isRunning(state: RunState): boolean {
257 return state.phase === 'compiling' || state.phase === 'loading' || state.phase === 'sampling' || state.phase === 'writing';
258 }
259}
261function el<K extends keyof HTMLElementTagNameMap>(tag: K, className?: string, text?: string): HTMLElementTagNameMap[K] {
262 const node = document.createElement(tag);
263 if (className) {
264 node.className = className;
265 }
266 if (text !== undefined) {
267 node.textContent = text;
268 }
269 return node;
270}
272function field(label: string, hint: string, control: HTMLElement, compact = false): HTMLElement {
273 const wrap = el('div', 'sample-field');
274 const labelEl = el('label', undefined, label);
275 if (!compact) {
276 labelEl.appendChild(el('span', 'sample-field-hint', hint));
277 } else {
278 labelEl.title = hint;
279 }
280 wrap.append(labelEl, control);
281 return wrap;
282}
284async function listProjectFiles(fs: WorkspaceFileSystem, path = '/'): Promise<string[]> {
285 const result: string[] = [];
286 const stat = await fs.fileService.resolve(fs.root.with({ path }));
287 for (const child of stat.children ?? []) {
288 if (child.isDirectory) {
289 result.push(...await listProjectFiles(fs, child.resource.path));
290 } else {
291 result.push(child.resource.path);
292 }
293 }
294 return result.sort();
295}