# /// script
# requires-python = ">=3.12"
# dependencies = ["marimo==0.25.2"]
# ///
# Generated by scripts/learn-notebooks/build.py; standalone and editable.
import marimo

__generated_with = "0.25.2"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo
    import math
    import struct
    from fractions import Fraction
    import json
    import inspect
    import hashlib
    from datetime import datetime, timezone
    from copy import deepcopy
    return Fraction, datetime, deepcopy, hashlib, inspect, json, math, mo, struct, timezone


@app.cell
def _():
    protocol = {'title': 'When does pooling preserve a shifted response?',
     'reference': {'row': 0, 'column': 0, 'method': 'max'},
     'controls': [{'key': 'row',
                   'label': 'Response row (0–3)',
                   'options': [{'value': 0, 'label': '0'},
                               {'value': 1, 'label': '1'},
                               {'value': 2, 'label': '2'},
                               {'value': 3, 'label': '3'}]},
                  {'key': 'column',
                   'label': 'Response column (0–3)',
                   'options': [{'value': 0, 'label': '0'},
                               {'value': 1, 'label': '1'},
                               {'value': 2, 'label': '2'},
                               {'value': 3, 'label': '3'}]},
                  {'key': 'method',
                   'label': 'Pooling operation',
                   'options': [{'value': 'max', 'label': 'Max pooling'},
                               {'value': 'mean', 'label': 'Average pooling'}]}],
     'fixed': 'One response of value 1 in a 4 × 4 feature map; every other cell is zero. Pool each '
              'non-overlapping 2 × 2 window with stride 2 and no padding. This isolates pooling, not a '
              'trained CNN.',
     'sources': [{'title': 'Understanding CNNs — convolutions, feature maps, and pooling',
                  'url': 'https://indepth.dev/posts/2007/en/understanding-cnns-convolutions-feature-maps-pooling'}],
     'id': 'pooling',
     'guided': 'Keep row = 0 and max pooling fixed. Run columns 0, 1, then 2. Compare every pooled '
               'entry, not only their sum.',
     'extension': 'Choose shifts within and across pooling windows, then test average pooling as a '
                  'separate comparison. Extend calculate with another response pattern or window '
                  'arrangement. Predict which translations preserve the complete output and seek a '
                  'counterexample.',
     'reasoning': 'Columns 0 and 1 belong to the same pooling window; column 2 crosses its boundary. '
                  'The pooled sum stays 1 for max pooling, even when the representation moves. An '
                  'unchanged sum does not establish translation invariance of a representation or '
                  'classifier.',
     'referenceDataset': {'id': 'pooling',
                          'version': 1,
                          'settings': {'row': 0, 'column': 0, 'method': 'max'},
                          'title': 'A sparse 4 × 4 feature map',
                          'unit': 'pixel',
                          'description': 'One response of value 1, with zero in every other position. '
                                         'Each table row is a pixel in one feature map, not a separate '
                                         'image.',
                          'columns': [{'key': 'id',
                                       'label': 'Row',
                                       'type': 'category',
                                       'role': 'identifier',
                                       'description': 'Stable identifier within this synthetic '
                                                      'fixture.'},
                                      {'key': 'row',
                                       'label': 'Row position',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Zero-based vertical coordinate.'},
                                      {'key': 'column',
                                       'label': 'Column position',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Zero-based horizontal coordinate.'},
                                      {'key': 'value',
                                       'label': 'Response',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'The input activation at this position.'}],
                          'rows': [{'id': '0,0', 'row': 0, 'column': 0, 'value': 1},
                                   {'id': '0,1', 'row': 0, 'column': 1, 'value': 0},
                                   {'id': '0,2', 'row': 0, 'column': 2, 'value': 0},
                                   {'id': '0,3', 'row': 0, 'column': 3, 'value': 0},
                                   {'id': '1,0', 'row': 1, 'column': 0, 'value': 0},
                                   {'id': '1,1', 'row': 1, 'column': 1, 'value': 0},
                                   {'id': '1,2', 'row': 1, 'column': 2, 'value': 0},
                                   {'id': '1,3', 'row': 1, 'column': 3, 'value': 0},
                                   {'id': '2,0', 'row': 2, 'column': 0, 'value': 0},
                                   {'id': '2,1', 'row': 2, 'column': 1, 'value': 0},
                                   {'id': '2,2', 'row': 2, 'column': 2, 'value': 0},
                                   {'id': '2,3', 'row': 2, 'column': 3, 'value': 0},
                                   {'id': '3,0', 'row': 3, 'column': 0, 'value': 0},
                                   {'id': '3,1', 'row': 3, 'column': 1, 'value': 0},
                                   {'id': '3,2', 'row': 3, 'column': 2, 'value': 0},
                                   {'id': '3,3', 'row': 3, 'column': 3, 'value': 0}],
                          'notes': ['This is one image-like fixture containing 16 pixels. It is not a '
                                    'dataset of 16 images.',
                                    '15 of 16 responses are zero. Inspect the spatial layout as well '
                                    'as the value distribution.',
                                    'Pooling type changes the operation, not the input map. Moving the '
                                    'response changes the input.']},
     'lessonUrl': 'https://indepth.dev/learn/experiments/machine-learning/pooling'}
    return protocol,


@app.cell(hide_code=True)
def _(mo, protocol):
    mo.md(f"""
    # {protocol['title']}

    [Read the lesson]({protocol['lessonUrl']}). This notebook is an editable Python companion.
    These small deterministic calculations run on CPU; a GPU is unnecessary.
    They do not reproduce full training runs from the source articles.

    **Protocol:** predict → run a reference → change one condition → compare → explain → design a variation.

    {protocol['fixed']}

    **First run:** use the reference settings shown below and write your prediction.
    **Then investigate:** {protocol['guided']}

    Runs remain in this running notebook session. Download your report before restarting or leaving.
    Editing the function does not recalculate old evidence; each new run saves the function source.
    """)
    return


@app.cell(hide_code=True)
def _(json, mo, protocol):
    _dataset = protocol['referenceDataset']
    mo.vstack([mo.md('## Inspect the reference inputs\n\n' + _dataset['description']),
               mo.ui.table(_dataset['rows'], selection=None, show_column_summaries=False),
               mo.md('\n'.join('- ' + note for note in _dataset['notes'])),
               mo.md('These are reference inputs. The input-dependent settings and full calculation source are saved with each run.'),
               mo.download(json.dumps(_dataset, indent=2).encode(), filename=protocol['id'] + '-reference-inputs.json', mimetype='application/json', label='Download reference inputs')])
    return


@app.cell
def _(Fraction, math, struct):
    def calculate(s):
        feature_map = [[0] * 4 for _ in range(4)]
        feature_map[s['row']][s['column']] = 1
        pooled = [[0.0] * 2 for _ in range(2)]
        for row in range(2):
            for col in range(2):
                window = [feature_map[2 * row + dr][2 * col + dc] for dr in range(2) for dc in range(2)]
                pooled[row][col] = max(window) if s['method'] == 'max' else sum(window) / 4
        return {'metrics': {'Top-left pooled value': pooled[0][0], 'Sum of pooled values': sum(map(sum, pooled))},
                'rows': [{'Pooled row': i, 'Column 0': row[0], 'Column 1': row[1]} for i, row in enumerate(pooled)],
                'input': feature_map, 'output': pooled, 'status': 'Completed'}
    return calculate,


@app.cell
def _(mo):
    get_runs, set_runs = mo.state([])
    return get_runs, set_runs


@app.cell
def _(calculate, datetime, deepcopy, hashlib, inspect, mo, protocol, set_runs, timezone):
    def record_run(value):
        if value is None or not value['prediction'].strip():
            return
        settings = {key: value[key] for key in protocol['reference']}
        source = inspect.getsource(calculate)
        snapshot = {
            'recordedAt': datetime.now(timezone.utc).isoformat(),
            'settings': settings,
            'prediction': value['prediction'].strip(),
            'implementationNote': value['implementation'],
            'calculationSource': source,
            'calculationSha256': hashlib.sha256(source.encode()).hexdigest(),
            'result': calculate(settings),
        }
        set_runs(lambda previous: [*previous, deepcopy(snapshot)])

    _controls = {}
    _parts = []
    for _control in protocol['controls']:
        _options = {choice['label']: choice['value'] for choice in _control['options']}
        _default = next(label for label, value in _options.items() if value == protocol['reference'][_control['key']])
        _controls[_control['key']] = mo.ui.dropdown(_options, value=_default, label=_control['label'])
        _parts.append('{' + _control['key'] + '}')
    _controls['prediction'] = mo.ui.text_area(label='Prediction and reason for this run', full_width=True)
    _controls['implementation'] = mo.ui.text_area(value='Original lesson calculation', label='Describe any code changes', full_width=True)
    _parts.extend(['{prediction}', '{implementation}'])
    experiment_form = mo.md('\n\n'.join(_parts)).batch(**_controls).form(
        submit_button_label='Run experiment',
        validate=lambda value: None if value and value['prediction'].strip() else 'Write a prediction and reason before running.',
        on_change=record_run,
    )
    experiment_form
    return experiment_form,


@app.cell(hide_code=True)
def _(get_runs, mo):
    runs = get_runs()
    baseline = mo.ui.dropdown({f"Run {i + 1}": i for i in range(len(runs))}, value='Run 1' if runs else None, label='Comparison baseline')
    mo.vstack([mo.md(f'## Recorded evidence — {len(runs)} runs'), baseline])
    return baseline, runs


@app.cell(hide_code=True)
def _(baseline, math, mo, runs):
    mo.stop(not runs, mo.md('Write a prediction, then select **Run experiment** to record the reference.'))
    _current = runs[-1]
    _base = runs[baseline.value if baseline.value is not None else 0]
    _changes = [key for key, value in _current['settings'].items() if value != _base['settings'][key]]
    _code_changed = _current['calculationSha256'] != _base['calculationSha256']
    _metrics = [{'Measurement': key, 'Baseline': _base['result']['metrics'].get(key), 'Latest run': value,
                 'Difference': value - _base['result']['metrics'][key] if key in _base['result']['metrics'] else None}
                for key, value in _current['result']['metrics'].items()]
    _outputs = [mo.md(f"**Latest run:** {len(runs)} · {_current['result']['status']}"),
                mo.md('**Prediction:** ' + _current['prediction']),
                mo.md('**Changed settings:** ' + (', '.join(_changes) or 'none') + '. **Calculation changed:** ' + str(_code_changed) + '.'),
                mo.md('Compare runs that differ in one intended condition. Check the saved code when comparing implementations.')]
    if _metrics:
        _outputs.append(mo.ui.table(_metrics, selection=None, show_column_summaries=False))
    if _current['result']['rows'] and 'Training loss' in _current['result']['rows'][0]:
        _training = [math.log10(1 + row['Training loss']) for row in _current['result']['rows']]
        _validation = [math.log10(1 + row['Validation loss']) for row in _current['result']['rows']]
        _top = max(1, max(_training), max(_validation))
        _lines = []
        for _values, _color, _dash in [(_training, '#2337ff', ''), (_validation, '#ae5a0b', '7 4')]:
            _points = ' '.join(f'{40 + 500 * i / max(1, len(_values) - 1)},{200 - 160 * value / _top}' for i, value in enumerate(_values))
            _lines.append(f'<polyline points="{_points}" fill="none" stroke="{_color}" stroke-dasharray="{_dash}" stroke-width="2"/>')
        _outputs.append(mo.Html(f'<svg viewBox="0 0 600 250" role="img" aria-label="Training and validation loss trajectories"><text x="40" y="20">log10(1 + loss), top = {_top:.3g}</text><path d="M40 40 V200 H540" fill="none" stroke="gray"/>{"".join(_lines)}<text x="40" y="225">0</text><text x="460" y="225">Update {len(_training)-1}</text></svg>'))
        _outputs.append(mo.md('Blue solid = training; orange dashed = validation. Use the table for exact losses.'))
    _outputs.append(mo.accordion({'Latest run: all intermediate evidence': mo.ui.table(_current['result']['rows'], selection=None, show_column_summaries=False)}))
    _outputs.append(mo.accordion({f"Run {i + 1}: settings, prediction, code, and evidence": mo.json(run) for i, run in enumerate(runs)}))
    mo.vstack(_outputs)
    return


@app.cell(hide_code=True)
def _(mo, protocol):
    observations = mo.ui.text_area(label='Observations: cite runs and measured values', full_width=True, debounce=False)
    explanation = mo.ui.text_area(label='Explanation: why did this happen?', full_width=True, debounce=False)
    limitations = mo.ui.text_area(label='Limitations: what remains unresolved?', full_width=True, debounce=False)
    next_test = mo.ui.text_area(label='Your next test and a result that would challenge your explanation', full_width=True, debounce=False)
    mo.vstack([mo.md('## Explain, then extend'), observations, explanation, limitations,
               mo.md('### Design your own investigation\n\n' + protocol['extension']), next_test,
               mo.accordion({'Check the reasoning after writing your explanation': mo.md(protocol['reasoning'])})])
    return explanation, limitations, next_test, observations


@app.cell(hide_code=True)
def _(baseline, datetime, explanation, json, limitations, mo, next_test, observations, protocol, runs, timezone):
    report = {'schemaVersion': 1, 'protocol': protocol, 'exportedAt': datetime.now(timezone.utc).isoformat(),
              'comparisonRunNumber': baseline.value + 1 if baseline.value is not None else None, 'runs': runs, 'notes': {'observations': observations.value, 'explanation': explanation.value,
                                     'limitations': limitations.value, 'nextTest': next_test.value}}
    mo.download(json.dumps(report, indent=2, allow_nan=False).encode(), filename=protocol['id'] + '-notebook-report.json',
                mimetype='application/json', label='Download report', disabled=not runs)
    return report,


if __name__ == '__main__':
    app.run()
