# /// 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': 'Can confidence rise without accuracy improving?',
     'reference': {'temperature': 1, 'offset': 0},
     'controls': [{'key': 'temperature',
                   'label': 'Temperature T',
                   'options': [{'value': 0.25, 'label': '0.25'},
                               {'value': 0.5, 'label': '0.5'},
                               {'value': 1, 'label': '1'},
                               {'value': 2, 'label': '2'},
                               {'value': 4, 'label': '4'}]},
                  {'key': 'offset',
                   'label': 'Common logit offset',
                   'options': [{'value': -1000, 'label': '-1000'},
                               {'value': -5, 'label': '-5'},
                               {'value': 0, 'label': '0'},
                               {'value': 5, 'label': '5'},
                               {'value': 1000, 'label': '1000'}]}],
     'fixed': 'Three fixed examples: logits (2,0,0), (2,0,0), (0,0,2); true classes 0,1,2. Apply the '
              'same positive temperature and common offset to each row. No training; cross-entropy '
              'uses natural logarithms and a stable log-sum-exp calculation.',
     'sources': [{'title': 'Training a real network on MNIST',
                  'url': 'https://indepth.dev/posts/1008/en/training-and-experimenting-with-mnist'},
                 {'title': 'The effect of hyperparameters on training',
                  'url': 'https://indepth.dev/posts/1012/en/effect-of-hyperparameters-on-training'}],
     'id': 'confidence',
     'guided': 'Keep the common offset at zero. Run temperatures 1, 0.5, and 0.25 separately. Compare '
               'accuracy, predicted-class probabilities, and each example’s loss.',
     'extension': 'Choose another positive temperature and explain which examples drive mean loss. '
                  'Test common offsets separately. Edit the logits and targets to create a contrasting '
                  'case. What additional data would you need before claiming improved calibration?',
     'reasoning': 'Positive temperature scaling preserves argmax, so these predictions and accuracy '
                  'stay fixed. Lower temperature sharpens both correct and incorrect predictions. A '
                  'confidently wrong example can dominate mean cross-entropy. These three frozen rows '
                  'are not a calibration study. Stable softmax also preserves probabilities under a '
                  'common offset.',
     'referenceDataset': {'id': 'confidence',
                          'version': 1,
                          'settings': {'temperature': 1, 'offset': 0},
                          'title': 'Frozen logits and known labels',
                          'unit': 'observation',
                          'description': 'Three constructed examples with fixed model scores. Two have '
                                         'identical logits but different true labels, letting us '
                                         'investigate confidence independently of accuracy.',
                          'columns': [{'key': 'id',
                                       'label': 'Row',
                                       'type': 'category',
                                       'role': 'identifier',
                                       'description': 'Stable identifier within this synthetic '
                                                      'fixture.'},
                                      {'key': 'logit0',
                                       'label': 'Logit 0',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Frozen score for class 0; not a probability.'},
                                      {'key': 'logit1',
                                       'label': 'Logit 1',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Frozen score for class 1; not a probability.'},
                                      {'key': 'logit2',
                                       'label': 'Logit 2',
                                       'type': 'number',
                                       'role': 'feature',
                                       'description': 'Frozen score for class 2; not a probability.'},
                                      {'key': 'target',
                                       'label': 'True class',
                                       'type': 'category',
                                       'role': 'target',
                                       'description': 'Known class label, independent of the model’s '
                                                      'score.'}],
                          'rows': [{'id': 'A', 'logit0': 2, 'logit1': 0, 'logit2': 0, 'target': 0},
                                   {'id': 'B', 'logit0': 2, 'logit1': 0, 'logit2': 0, 'target': 1},
                                   {'id': 'C', 'logit0': 0, 'logit1': 0, 'logit2': 2, 'target': 2}],
                          'notes': ['A and B share the same features but have different labels. They '
                                    'are feature duplicates, not duplicate complete rows; do not '
                                    'automatically remove either.',
                                    'These three examples cannot establish population calibration or '
                                    'model generalization.',
                                    'Temperature and common offset transform these frozen scores '
                                    'during the experiment; the source logits remain unchanged.']},
     'lessonUrl': 'https://indepth.dev/learn/experiments/machine-learning/confidence'}
    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):
        inputs, targets = [(2, 0, 0), (2, 0, 0), (0, 0, 2)], [0, 1, 2]
        rows, loss, confidence_sum, correct = [], 0.0, 0.0, 0
        for i, (logits, target) in enumerate(zip(inputs, targets)):
            scaled = [(value + s['offset']) / s['temperature'] for value in logits]
            largest = max(scaled)
            exp = [math.exp(value - largest) for value in scaled]
            denominator = sum(exp)
            probabilities = [value / denominator for value in exp]
            prediction = scaled.index(largest)
            cross_entropy = largest - scaled[target] + math.log(denominator)
            correct += prediction == target
            loss += cross_entropy
            confidence_sum += max(probabilities)
            rows.append({'Example': chr(65 + i), 'True class': target, 'Predicted': prediction, **{f'P(class {j})': p for j, p in enumerate(probabilities)}, 'Cross-entropy': cross_entropy})
        return {'metrics': {'Accuracy (%)': 100 * correct / 3, 'Mean predicted-class probability': confidence_sum / 3, 'Mean cross-entropy': loss / 3}, 'rows': rows, '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()
