# /// 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': 'Do all paths contribute to the gradient?',
     'reference': {'x': 2, 'y': 3, 'paths': 'both', 'h': 0.001},
     'controls': [{'key': 'x',
                   'label': 'Input x',
                   'options': [{'value': -3, 'label': '-3'},
                               {'value': -2, 'label': '-2'},
                               {'value': -1, 'label': '-1'},
                               {'value': 0, 'label': '0'},
                               {'value': 1, 'label': '1'},
                               {'value': 2, 'label': '2'},
                               {'value': 3, 'label': '3'}]},
                  {'key': 'y',
                   'label': 'Fixed input y',
                   'options': [{'value': -3, 'label': '-3'},
                               {'value': -2, 'label': '-2'},
                               {'value': -1, 'label': '-1'},
                               {'value': 0, 'label': '0'},
                               {'value': 1, 'label': '1'},
                               {'value': 2, 'label': '2'},
                               {'value': 3, 'label': '3'}]},
                  {'key': 'paths',
                   'label': 'Backward paths',
                   'options': [{'value': 'both', 'label': 'Both paths'},
                               {'value': 'direct', 'label': 'Direct multiplication path only'},
                               {'value': 'sum', 'label': 'Path through the sum only'}]},
                  {'key': 'h',
                   'label': 'Difference step h',
                   'options': [{'value': 0.1, 'label': '0.1'},
                               {'value': 0.001, 'label': '0.001'},
                               {'value': 1e-06, 'label': '0.000001'},
                               {'value': 1e-12, 'label': '1e-12'},
                               {'value': 1e-16, 'label': '1e-16'}]}],
     'fixed': 'Forward computation: a = x + y, z = a × x. The backward-path setting changes only the '
              'reported gradient. Central differences always evaluate the complete forward function '
              'with y fixed.',
     'sources': [{'title': 'From derivatives to backpropagation',
                  'url': 'https://indepth.dev/posts/2124/en/from-derivatives-to-backpropagation'}],
     'id': 'gradient-paths',
     'guided': 'Run both paths at x = 2, y = 3, h = 0.001. Then omit each path separately, keeping the '
               'forward function and inputs fixed. Compare reported derivatives with central '
               'differences.',
     'extension': 'Find an input that hides an omitted contribution, then one that exposes it. Explain '
                  'why one successful gradient check is insufficient. Separately vary h to explore '
                  'cancellation. For a code extension, replace the forward computation and derive '
                  'every backward path before testing it.',
     'reasoning': 'The derivative is 2x + y: the direct contribution is x + y and the sum path '
                  'contributes x. At x = 0 the omitted sum contribution vanishes; this does not '
                  'validate an incomplete rule. Very small finite-difference steps can lose '
                  'information to rounding.',
     'referenceDataset': {'id': 'gradient-paths',
                          'version': 1,
                          'settings': {'x': 2, 'y': 3, 'paths': 'both', 'h': 0.001},
                          'title': 'Inputs for a numerical gradient check',
                          'unit': 'evaluation point',
                          'description': 'These are the center and two nearby inputs used to evaluate '
                                         'z = (x + y)x. They make the finite-difference comparison '
                                         'inspectable before calculating the gradient.',
                          'columns': [{'key': 'id',
                                       'label': 'Row',
                                       'type': 'category',
                                       'role': 'identifier',
                                       'description': 'Stable identifier within this synthetic '
                                                      'fixture.'},
                                      {'key': 'probe',
                                       'label': 'Evaluation',
                                       'type': 'category',
                                       'role': 'condition',
                                       'description': 'Center, left, or right probe; these are '
                                                      'dependent evaluation points.'},
                                      {'key': 'x',
                                       'label': 'Input x',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'The varied input, stored as a float64 number.'},
                                      {'key': 'y',
                                       'label': 'Input y',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Held fixed during the gradient check.'}],
                          'rows': [{'id': 'center', 'probe': 'Center', 'x': 2, 'y': 3},
                                   {'id': 'left', 'probe': 'Left', 'x': 1.999, 'y': 3},
                                   {'id': 'right', 'probe': 'Right', 'x': 2.001, 'y': 3}],
                          'notes': ['Three deterministic probes, not three independent training '
                                    'samples.',
                                    'At very small h, float64 may store a perturbed x as the same '
                                    'value as the center. Feature duplicates can reveal that loss of '
                                    'resolution.',
                                    'Changing the backward-path setting leaves these inputs and the '
                                    'forward function unchanged.']},
     'lessonUrl': 'https://indepth.dev/learn/experiments/machine-learning/gradient-paths'}
    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):
        x, y, h = s['x'], s['y'], s['h']
        direct, via_sum = x + y, x
        gradient = direct + via_sum if s['paths'] == 'both' else direct if s['paths'] == 'direct' else via_sum
        def forward(value):
            return (value + y) * value
        estimate = (forward(x + h) - forward(x - h)) / (2 * h)
        return {'metrics': {'Forward output z': forward(x), 'Reported gradient': gradient, 'Central difference': estimate, 'Absolute discrepancy': abs(gradient - estimate)},
                'rows': [{'Path': 'Direct multiplication input', 'Local product': direct, 'Included contribution': 0 if s['paths'] == 'sum' else direct},
                         {'Path': 'Through a = x + y', 'Local product': via_sum, 'Included contribution': 0 if s['paths'] == 'direct' else via_sum}], '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()
