# /// script
# requires-python = ">=3.12"
# dependencies = ["marimo==0.24.0", "numpy>=2.0,<3"]
# ///
import marimo
__generated_with = "0.24.0"
app = marimo.App(width="full")

@app.cell
def _():
    """Spark experiments. Small-model extensions, not reproductions of Goodfire's models."""
    import numpy as np
    import json, re, html, sys, os, importlib
    SMOL='HuggingFaceTB/SmolLM2-135M-Instruct'
    SMOL_REV='12fd25f77366fa6b3b4b768ec3050bf629380bac'
    MINI='sentence-transformers/all-MiniLM-L6-v2'
    MINI_REV='1110a243fdf4706b3f48f1d95db1a4f5529b4d41'
    ORIGINAL='Pip brought a paper boat to the creek. Wattle said the water looked too fast. Pip tried anyway, and the boat disappeared around a rock. They walked downstream together. The boat was caught safely in some reeds. Pip thanked Wattle and decided to test the next boat in a puddle.'
    EMOTIONS=['surprise','disgust','anger','happiness','sadness','fear']
    MONTHS=['January','February','March','April','May','June','July','August','September','October','November','December']
    DAYS=['Monday','Tuesday','Wednesday','Thursday','Friday','Saturday','Sunday']

    def sm(x):
        e=np.exp(x-np.max(x,axis=-1,keepdims=True));return e/e.sum(-1,keepdims=True)

    def projection(ref,other=None):
        ref=np.asarray(ref,dtype=float);mean=ref.mean(0);centre=ref-mean
        _,s,v=np.linalg.svd(centre,full_matrices=False);basis=v[:2]
        return (np.asarray(ref if other is None else other)-mean)@basis.T,float((s[:2]**2).sum()/max((s**2).sum(),1e-20))

    def plot(points, labels=None, second=None, axes=('PC1','PC2'), connect=True):
        points=np.asarray(points);second=np.asarray(second) if second is not None else None
        allp=points if second is None else np.vstack([points,second]);scale=max(.001,float(np.abs(allp).max()))
        def xy(p):return 300+245*p[0]/scale,200-150*p[1]/scale
        out=['<svg viewBox="0 0 600 400" role="img" aria-label="Two-dimensional representation plot" style="max-width:100%;background:#f3f7fc"><path d="M40 200H560 M300 35V365" stroke="#bcc8cd"/><text x="525" y="225">PC1</text><text x="310" y="35">PC2</text>']
        for n,arr in enumerate([points] if second is None else [points,second]):
            colour=['#195cca','#a23c76'][n];coords=[xy(p) for p in arr]
            if connect: out.append('<polyline fill="none" stroke="'+colour+'" stroke-width="2" '+('stroke-dasharray="6 4"' if n else '')+' points="'+' '.join(f'{x:.2f},{y:.2f}' for x,y in coords)+'"/>')
            for i,(x,y) in enumerate(coords):
                lab=str(i+1) if labels is None else str(labels[i]);out.append(f'<circle cx="{x:.2f}" cy="{y:.2f}" r="4" fill="{colour}"/><text x="{x+6:.2f}" y="{y-7:.2f}" font-size="12">{html.escape(lab)}</text>')
        return (''.join(out)+'</svg>').replace('>PC1<','>'+html.escape(axes[0])+'<').replace('>PC2<','>'+html.escape(axes[1])+'<')

    def train_addition(epochs=600,seed=7):
        """Train an actual one-hidden-layer network on 80 of 100 mod-10 sums."""
        rng=np.random.default_rng(seed);pairs=np.array([(a,b) for a in range(10) for b in range(10)])
        x=np.concatenate([np.eye(10)[pairs[:,0]],np.eye(10)[pairs[:,1]]],axis=1);y=pairs.sum(1)%10
        split=rng.permutation(100);train,test=split[:80],split[80:]
        w=rng.normal(0,.25,(20,48));b=np.zeros(48);v=rng.normal(0,.15,(48,10));c=np.zeros(10);history=[]
        for epoch in range(int(epochs)):
            h=np.maximum(0,x[train]@w+b);p=sm(h@v+c);d=(p-np.eye(10)[y[train]])/len(train)
            dh=(d@v.T)*(h>0);w-=.6*(x[train].T@dh+1e-4*w);b-=.6*dh.sum(0);v-=.6*(h.T@d+1e-4*v);c-=.6*d.sum(0)
            if epoch%100==0:history.append({'step':epoch,'training_loss':float(-np.log(p[np.arange(80),y[train]]+1e-12).mean())})
        hidden=np.maximum(0,x@w+b);pred=(hidden@v+c).argmax(1)
        # Class means are descriptive summaries, not a causal intervention.
        means=np.array([hidden[y==k].mean(0) for k in range(10)]);points,var=projection(means)
        return {'model':'trained 20 → 48 ReLU → 10 network','seed':seed,'epochs':epochs,'train_indices':train.tolist(),'test_indices':test.tolist(),'train_correct':int((pred[train]==y[train]).sum()),'train_total':80,'test_correct':int((pred[test]==y[test]).sum()),'test_total':20,'pca_variance':var,'rows':[{'a':int(pairs[i,0]),'b':int(pairs[i,1]),'true_mod10_sum':int(y[i]),'prediction':int(pred[i]),'split':'test' if i in test else 'train'} for i in range(100)],'history':history,'svg':plot(points,list(range(10))),'interpretation':'The plotted points are learned hidden-state means grouped by known sum. They need not form a circle. Grouping itself uses labels. A hand-written Fourier circle is a separate mathematical construction.'}

    _NATIVE={}
    def language_model():
        if sys.platform=='emscripten':raise RuntimeError('This generative-model extension needs native Python. Download this notebook and follow the Spark setup guide. The browser activities above still work.')
        if 'smol' not in _NATIVE:
            torch=importlib.import_module('torch')
            transformers=importlib.import_module('transformers')
            AutoTokenizer=transformers.AutoTokenizer;AutoModelForCausalLM=transformers.AutoModelForCausalLM
            tok=AutoTokenizer.from_pretrained(SMOL,revision=SMOL_REV)
            device=os.environ.get('SPARK_DEVICE','cpu')
            if device not in ['cpu','cuda']:raise ValueError('SPARK_DEVICE must be cpu or cuda.')
            if device=='cuda' and not torch.cuda.is_available():raise RuntimeError('CUDA is unavailable. Use CPU or a configured GPU runtime.')
            model=AutoModelForCausalLM.from_pretrained(SMOL,revision=SMOL_REV).eval().to(device)
            _NATIVE['smol']=(tok,model)
        return _NATIVE['smol']

    def hidden_and_logits(text,layer=15,max_tokens=512):
        torch=importlib.import_module('torch')
        tok,model=language_model();ids=tok(text,return_tensors='pt',add_special_tokens=False).input_ids;count=ids.shape[1];ids=ids[:,-max_tokens:].to(model.device)
        with torch.no_grad():out=model(ids,output_hidden_states=True)
        layer=max(0,min(int(layer),len(out.hidden_states)-1))
        return out.hidden_states[layer][0,-1].float().cpu().numpy(),out.logits[0,-1].float().cpu().numpy(),int(count),layer

    def native_calculator(layer=15,steer_to=17,seed=7,custom_domain="months",custom_start=8,custom_offset=6):
        tok,model=language_model();pairs=[(a,b) for a in range(1,9) for b in range(1,9)]
        chat=lambda p:tok.apply_chat_template([{'role':'user','content':p}],tokenize=False,add_generation_prompt=True)
        prompts=[chat(f'Calculate {a} + {b}. Give only the number.') for a,b in pairs];sums=np.array([a+b for a,b in pairs])
        states=[];rows=[]
        for prompt,total in zip(prompts,sums):
            h,logits,_,actual_layer=hidden_and_logits(prompt,layer);states.append(h)
            # Full next-token prediction, not a forced choice among numbers.
            rows.append({'prompt':prompt,'sum':int(total),'next_token':tok.decode([int(logits.argmax())])})
        x=np.array(states);rng=np.random.default_rng(seed);idx=rng.permutation(64);train,test=idx[:48],idx[48:]
        mu=x[train].mean(0);sd=x[train].std(0)+1e-3;z=(x-mu)/sd
        def fourier(n):return np.array([f(2*np.pi*n/p) for p in [2,5,10] for f in [np.cos,np.sin]])
        targets=np.array([fourier(n) for n in sums]);xt=np.column_stack([z[train],np.ones(48)])
        # Dual ridge avoids a large matrix inverse; only training examples fit the probe.
        w=xt.T@np.linalg.solve(xt@xt.T+5*np.eye(48),targets[train]);pred=np.column_stack([z,np.ones(64)])@w
        baseline=targets[train].mean(0);err=float(np.sqrt(((pred[test]-targets[test])**2).mean()));null=float(np.sqrt(((baseline-targets[test])**2).mean()))
        time_cases=[('August + 6 months','What month is 6 months after August? Answer:',14,'February'),('August + 16 months','What month is 16 months after August? Answer:',24,'December'),('Friday + 2 days','What day is 2 days after Friday? Answer:',7,'Sunday'),('13 + 4 hours','What hour is 4 hours after 13:00? Answer:',17,'17'),('23 + 3 hours','What hour is 3 hours after 23:00? Answer:',26,'2')]
        if custom_domain=='months':
            if not 1<=custom_start<=12:raise ValueError('Months use start 1–12.')
            custom=f'What month is {custom_offset} months after {MONTHS[custom_start-1]}?';expected=MONTHS[(custom_start+custom_offset-1)%12]
        elif custom_domain=='weekdays':
            if not 1<=custom_start<=7:raise ValueError('Weekdays use Monday=1 through Sunday=7.')
            custom=f'What day is {custom_offset} days after {DAYS[custom_start-1]}?';expected=DAYS[(custom_start+custom_offset-1)%7]
        else:
            if not 0<=custom_start<=23:raise ValueError('Hours use 0–23.')
            custom=f'What hour is {custom_offset} hours after {custom_start}:00?';expected=str((custom_start+custom_offset)%24)
        time_cases.append(('Your time example',custom,custom_start+custom_offset,expected))
        def generate_short(prompt):
            torch=importlib.import_module('torch')
            ids=tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device)
            with torch.no_grad():out=model.generate(ids,max_new_tokens=40,do_sample=False,pad_token_id=tok.eos_token_id,attention_mask=torch.ones_like(ids))
            return tok.decode(out[0,ids.shape[1]:],skip_special_tokens=True)
        transfer=[]
        for name,prompt,total,expected in time_cases:
            h,logits,_,_=hidden_and_logits(chat(prompt),layer);p=np.append((h-mu)/sd,1)@w
            transfer.append({'task':name,'prompt':prompt,'ordinary_sum':total,'expected_answer':expected,'next_token':tok.decode([int(logits.argmax())]),'generated_answer':generate_short(chat(prompt)),'probe_rmse':float(np.sqrt(((p-fourier(total))**2).mean()))})
        # Intervene on post-block state using inverse of the fitted probe. Includes exact no-op control.
        torch=importlib.import_module('torch')
        prompt=chat('Calculate 7 + 9. Give only the number.');h,base,_,_=hidden_and_logits(prompt,layer)
        current=np.append((h-mu)/sd,1)@w;delta=(fourier(steer_to)-current)@np.linalg.pinv(w[:-1]);delta=delta*sd
        norm=float(np.linalg.norm(delta));cap=.15*float(np.linalg.norm(h));delta=delta*min(1,cap/max(norm,1e-12))
        def intervene(vector):
            def hook(_module,_inputs,output):
                value=output[0] if isinstance(output,tuple) else output
                changed=value.clone();changed[0,-1]+=torch.tensor(vector,dtype=changed.dtype,device=changed.device)
                return (changed,)+output[1:] if isinstance(output,tuple) else changed
            # hidden_states index k is post block k-1; final index includes final norm, so avoid it.
            handle=model.model.layers[actual_layer-1].register_forward_hook(hook)
            try:
                with torch.no_grad():out=model(tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device))
                return out.logits[0,-1].float().cpu().numpy()
            finally:handle.remove()
        if not 1<=actual_layer<len(model.model.layers):raise ValueError('Choose an intermediate layer from 1 to 29; final-normalised states are not valid for this intervention.')
        noop=intervene(np.zeros_like(delta));changed=intervene(delta)
        return {'model':SMOL,'revision':SMOL_REV,'layer':actual_layer,'device':str(model.device),'seed':seed,'training_indices':train.tolist(),'test_indices':test.tolist(),'heldout_probe_rmse':err,'constant_baseline_rmse':null,'next_token_rows':rows,'rows':transfer,'steering':{'prompt':prompt,'requested_sum':steer_to,'delta_norm':float(np.linalg.norm(delta)),'base_next_token':tok.decode([int(base.argmax())]),'steered_next_token':tok.decode([int(changed.argmax())]),'noop_max_logit_error':float(np.abs(base-noop).max()),'mean_absolute_logit_change':float(np.abs(changed-base).mean())},'svg':plot(pred[test,4:6],[int(sums[i]) for i in test],axes=('cos(10)','sin(10)'),connect=False),'plot_axes':'Predicted cos/sin for period 10; Points are held-out prompts, not a trajectory.','interpretation':'Poor probe performance or transfer is evidence against this classroom hypothesis for this model/prompt/layer. Next-token fragments are not scored as full answers. Steering can disrupt computation; a changed answer is not proof of the complete Llama algorithm.'}

    def truth_experiment(steps=200,seed=7):
        # MODEL and forward are injected from the independently trained Brightlab transformer.
        rng=np.random.default_rng(seed);cols=[MODEL['vocab'].index(c) for c in ['red','blue','green','gold']]
        features=[];truth=[];logits=[]
        for tokens,y in MODEL['cases']:
            l,c=forward(tokens,'edited');features.append(c['residual'][-1][-1]);logits.append(l[cols]);truth.append(cols.index(y))
        h=np.array(features);truth=np.array(truth);logits=np.array(logits);test=np.array(MODEL['test_indices']);train=np.array([i for i in range(len(h)) if i not in test])
        mu=h[train].mean(0);sd=h[train].std(0)+1e-3;x=np.column_stack([(h-mu)/sd,np.ones(len(h))])
        # One logistic classifier per proposed colour. It sees frozen internal features, not truth at test time.
        w=rng.normal(0,.01,(x.shape[1],4));target=np.eye(4)[truth]
        for _ in range(300):
            p=1/(1+np.exp(-np.clip(x[train]@w,-30,30)));w-=.08*(x[train].T@(p-target[train])/len(train)+.005*w)
        scores=1/(1+np.exp(-np.clip(x@w,-30,30)));reward=np.column_stack([scores,np.full(len(x),.55)])
        prior=np.column_stack([logits/3,np.max(logits/3,1)-1]);policy=np.zeros((x.shape[1],5));ref=sm(prior)
        for _ in range(int(steps)):
            p=sm(prior[train]+x[train]@policy);adv=reward[train]-.12*(np.log(p+1e-12)-np.log(ref[train]+1e-12));grad=p*(adv-(p*adv).sum(1,keepdims=True));policy+=.2*x[train].T@grad/len(train)
        after=sm(prior+x@policy)
        def metrics(p):
            a=p[test].argmax(1);answered=a!=4;wrong=answered&(a!=truth[test]);return {'questions':len(test),'answered':int(answered.sum()),'wrong_answers':int(wrong.sum()),'correct_answers':int((answered&(a==truth[test])).sum()),'wrong_among_answered':float(wrong.sum()/max(1,answered.sum())),'coverage':float(answered.mean())}
        names=['red','blue','green','gold','I need to check'];rows=[]
        for i in test:
            tokens,_=MODEL['cases'][i];chosen=int(after[i].argmax());rows.append({'case':int(i),'facts':' '.join(MODEL['vocab'][t] for t in tokens),'supported':names[truth[i]],'before':names[int(ref[i].argmax())],'after':names[chosen],'checker_for_after':None if chosen==4 else round(float(scores[i,chosen]),3)})
        return {'model':'Brightlab 10,522-parameter trained transformer, edited checkpoint; frozen hidden-state checker; separately reward-trained 25 × 5 answer policy','seed':seed,'steps':steps,'train_cases':len(train),'reserved_cases':len(test),'checker_accuracy_all_candidate_claims':float(((scores[test]>=.5)==target[test]).mean()),'before':metrics(ref),'after':metrics(after),'rows':rows,'reward':'Learned support score for colour answers; 0.55 for asking to check; KL penalty 0.12 against initial policy. Weights of transformer and checker stay frozen during policy training.','interpretation':'This is a small contextual-bandit policy update on known fictional facts, not full RLFR. A high checker score can be wrong. Ground truth is used to evaluate reserved cases, never as their training reward.'}

    def sentences(text):
        if len(text)>24000:raise ValueError('Select a passage shorter than 24,000 characters.')
        parts=[s.strip() for s in re.findall(r'[^.!?]+[.!?]+[”\"\x27]*|[^.!?]+$',text) if s.strip()]
        if not 3<=len(parts)<=40:raise ValueError('Use 3–40 sentences. Select a shorter passage from longer works.')
        return parts

    def story_inputs(text,ending,context):
        parts=sentences(text);variant=parts[:-1]+[ending.strip() or parts[-1]]
        def inputs(s):return [' '.join(s[:i+1]) if context=='accumulated' else s[i] for i in range(len(s))]
        return parts,variant,inputs(parts)+inputs(variant)

    def story_result(parts,variant,vectors,counts,model,revision,context,limit,pooling,emotion_rows=None):
        n=len(parts);vectors=np.array(vectors);points,var=projection(vectors[:n],vectors);rows=[]
        for i in range(n):
            distance=float(np.linalg.norm(vectors[i]-vectors[n+i]));cos=float(vectors[i]@vectors[n+i]/max(1e-12,np.linalg.norm(vectors[i])*np.linalg.norm(vectors[n+i])))
            rows.append({'sentence':i+1,'original':parts[i],'changed':variant[i],'PC1':float(points[i,0]),'PC2':float(points[i,1]),'changed_PC1':float(points[n+i,0]),'changed_PC2':float(points[n+i,1]),'tokens':counts[i],'changed_tokens':counts[n+i],'truncated':counts[i]>limit or counts[n+i]>limit,'vector_distance':distance,'cosine_similarity':cos})
        return {'model':model,'revision':revision,'context':context,'token_limit':limit,'pooling':pooling,'pca_variance':var,'projection':'PCA fitted to original version only; unchanged basis and mean for variant. Arbitrary component signs. Two axes are not inherently emotions.','rows':rows,'emotion_readouts':emotion_rows,'svg':plot(points[:n],second=points[n:]),'interpretation':'Blue solid = original; pink dashed = replacement ending. A projection can hide differences. Compare full-vector distances. Movement can reflect length, vocabulary, context and truncation as well as narrative changes.'}

    async def browser_story(text,ending,context='accumulated'):
        parts,variant,inputs=story_inputs(text,ending,context)
        globalThis=importlib.import_module('js').globalThis
        # Marimo runs Pyodide in a worker, which has no window. Import the inference module in that worker.
        origin=str(globalThis.location.origin)
        if origin=='null':raise RuntimeError('Open this notebook from the BrightLab site, not a local file URL, or use native mode.')
        module_url=origin+'/spark/model-runtime.js'
        program='(async()=>{if(!globalThis.sparkEmbed){await import('+json.dumps(module_url)+');}return await globalThis.sparkEmbed('+json.dumps(json.dumps(inputs))+');})()'
        data=json.loads(await globalThis.eval(program))
        return story_result(parts,variant,data['vectors'],data['tokenCounts'],data['model'],data['revision'],context,256,'MiniLM mean-pooled encoder; FIRST 256 tokens retained')

    def native_story(text,ending,context='accumulated',layer=15,emotions=False):
        parts,variant,inputs=story_inputs(text,ending,context);vectors=[];counts=[];emotion_rows=[]
        for t in inputs:
            h,_,count,actual=hidden_and_logits(t,layer,512);vectors.append(h);counts.append(count)
        if emotions:
            # Separate prompted behaviour: average token log likelihood of candidate labels, not 0–10 ratings.
            torch=importlib.import_module('torch')
            tok,model=language_model()
            for i,t in enumerate(inputs[:len(parts)]):
                prompt=f'Story: {t}\nThe main emotion in this story is';ids=tok(prompt,return_tensors='pt',add_special_tokens=False).input_ids[:,-480:].to(model.device);scores=[]
                for emotion in EMOTIONS:
                    suffix=tok(' '+emotion,return_tensors='pt',add_special_tokens=False).input_ids.to(model.device);full=torch.cat([ids,suffix],1)
                    with torch.no_grad():logp=model(full).logits[0].log_softmax(-1)
                    values=[float(logp[ids.shape[1]-1+j,token]) for j,token in enumerate(suffix[0])];scores.append(float(np.mean(values)))
                probs=sm(np.array(scores));emotion_rows.append({'sentence':i+1,**{e:float(p) for e,p in zip(EMOTIONS,probs)}})
        return story_result(parts,variant,vectors,counts,SMOL,SMOL_REV,context,512,f'Last-token hidden state at layer {actual}; LAST 512 tokens retained',emotion_rows or None)

    def fetch_book(book):
        # A curated identifier, not an arbitrary URL / private network fetch.
        if str(book) not in ['1524','2591']:raise ValueError('Choose Hamlet (1524) or Grimms’ Fairy Tales (2591).')
        url=f'https://www.gutenberg.org/ebooks/{book}.txt.utf-8'
        if sys.platform=='emscripten':raise RuntimeError('Gutenberg importing uses native Python. In the browser, paste a selected passage or upload a .txt file instead.')
        from urllib.request import urlopen,Request
        with urlopen(Request(url,headers={'User-Agent':'BrightLabSpark classroom research'}),timeout=45) as response:
            data=response.read(3_000_001)
        if len(data)>3_000_000:raise ValueError('This edition exceeds the 3 MB import limit.')
        return {'source':url,'book':str(book),'text':data.decode('utf-8-sig').replace('\r\n','\n').replace('\r','\n'),'note':'Choose and preview a 3–40-sentence passage. A full play is not one model context. The Gutenberg edition includes licence and editorial material.'}

    import marimo as mo
    lesson = {'id': 'y5-stories', 'year': 5, 'kind': 'stories', 'title': 'A story leaves a trail', 'question': 'How does a model’s picture of a story change as it reads another sentence?', 'goal': 'Connect a change in a story’s evidence with a change in a model’s representation; recognise that a model’s map is not a character’s feelings or a reader’s judgement.', 'prerequisite': 'Read aloud, notice a change in a character’s situation, and compare two points on a map.', 'source': 'https://www.goodfire.com/research/stories-in-space', 'paper': 'https://arxiv.org/html/2605.12412v1', 'minutes': 45, 'basis': 'The authors compare sentence-by-sentence emotion readouts with internal Llama representations. Their examples show trajectories as story context accumulates.', 'sequence': ['With your teacher, pick one of the stories in the authors’ original reader and predict a turning point.', 'Read one sentence at a time. Compare your interpretation with the model’s six emotion ratings.', 'Follow the same story in the original hidden-state map. A moving point represents changing numbers, not an AI feeling an emotion.', 'Paste a classroom story into the Spark notebook. Compare the original with a changed ending using a real sentence encoder.', 'For a teacher-led extension, run the notebook natively to collect actual hidden states from a small generative language model.'], 'assessment': 'Point to two sentences, describe the movement between them, and explain one reason your interpretation could differ from the model’s.', 'boundary': 'The source demos show the authors’ Llama experiments. The browser analyser uses MiniLM, a smaller sentence encoder; its map is a classroom extension. Native mode uses a generative language model. Neither proves a unique ‘shape’ for a story.'}
    MODEL = None
    forward = None
    engine = {'train_addition':train_addition,'native_calculator':native_calculator,'truth_experiment':truth_experiment,'browser_story':browser_story,'native_story':native_story,'fetch_book':fetch_book,'ORIGINAL':ORIGINAL}
    return (mo, json, sys, engine, lesson,)

@app.cell
def _(mo, lesson):
    mo.md(f"""# BrightLab Spark · Year {lesson['year']}
    ## {lesson['title']}
    **The question:** {lesson['question']}

    **Your goal:** {lesson['goal']}

    [Read Goodfire’s article]({lesson['source']}) · [Research paper]({lesson['paper']}) · [Spark setup guide](https://brightlab-ai-creators.ian347727.chatgpt.site/spark/field-guide)

    **Research boundary:** {lesson['boundary']}

    Work in pairs: predict first, run once, then use a control to challenge your explanation. Results and text stay in this session unless you download them.
    """)
    return 

@app.cell
def _(mo):
    prediction = mo.ui.text_area(label="Before running: what do you predict, and what evidence would change your mind?", full_width=True)
    prediction
    return (prediction,)

@app.cell
def _(mo):
    get_result, set_result = mo.state(None)
    return (get_result, set_result,)

@app.cell
def _(mo, sys):
    mo.md("""### Read → predict → map → change one thing
    Choose a short classroom story with **3–40 sentences**. Predict its turning point. Replace only the last sentence, then compare both versions on the **same map**.

    **Browser mode:** a real MiniLM sentence encoder downloads about 25 MB. It averages token vectors and keeps the first 256 tokens of each input. Long prefixes can lose the ending, so check the truncation column.

    **Native mode:** SmolLM2-135M reads each prefix; we collect its last-token hidden state **without asking an emotion question**. It keeps the last 512 tokens. This mode needs the downloaded notebook and native setup.

    Blue solid is the original; pink dashed is the changed ending. PCA makes a two-dimensional view of much longer vectors. Its axes do not mean “happy” and “sad”. These models do not feel the story.
    """)
    return 

@app.cell
def _():
    book_passage = ""
    book_source = "Classroom-authored passage"
    return (book_passage, book_source,)

@app.cell
def _(mo, engine, book_passage, book_source):
    story_form = mo.ui.dictionary({
    "text":mo.ui.text_area(value=book_passage or engine['ORIGINAL'],label="Preview and edit your original story (3–40 sentences)",rows=8,full_width=True),
    "upload":mo.ui.file(filetypes=['.txt'],multiple=False,label="Or upload a UTF-8 .txt passage (replaces editor text; maximum 24 KB)"),
    "ending":mo.ui.text_area(value="Pip blamed Wattle and left the boat behind.",label="Replacement final sentence",full_width=True),
    "context":mo.ui.dropdown(options=['accumulated','isolated'],value='accumulated',label="Read all sentences so far, or each sentence alone?"),
    "runtime":mo.ui.dropdown(options=['browser MiniLM','native SmolLM'],value='browser MiniLM',label="Model runtime"),
    "layer":mo.ui.number(start=1,stop=29,value=15,label="Native model layer"),
    "emotions":mo.ui.checkbox(value=False,label="Native: also score six candidate emotion labels (separate prompted readout)"),
    "source":mo.ui.text(value=book_source,label="Record the author, edition / source, and passage boundaries",full_width=True)
    }).form(submit_button_label="Analyse original and changed ending")
    story_form
    return (story_form,)

@app.cell
async def _(story_form, engine, set_result, sys):
    if story_form.value is not None:
        try:
            _s=dict(story_form.value);_uploads=_s.pop('upload');_text=_s['text']
            if _uploads:
                if len(_uploads[0].contents)>24000: raise ValueError("Upload a selected passage under 24 KB, not a whole book. Or import a Gutenberg edition natively and select paragraphs.")
                _text=_uploads[0].contents.decode('utf-8-sig');_s['source']+=' · uploaded '+_uploads[0].name
            _s['text']=_text
            if _s['runtime']=='browser MiniLM':
                if sys.platform!='emscripten': raise ValueError("Choose native SmolLM in the downloaded notebook. Browser MiniLM runs on the hosted page.")
                _data=await engine['browser_story'](_text,_s['ending'],_s['context'])
            else: _data=engine['native_story'](_text,_s['ending'],_s['context'],_s['layer'],_s['emotions'])
            set_result({"experiment":"Paired story representations","settings":_s,"result":_data})
        except Exception as _e: set_result({"error":str(_e)})
    return 

@app.cell
def _(get_result, mo, json):
    latest = get_result()
    if latest is None:
        display = mo.md("### Your evidence will appear here\nSubmit an experiment above. Model calculations only start when you submit.")
    elif 'error' in latest:
        display = mo.callout(latest['error'],kind="warn")
    else:
        _r=latest['result'];_overview={k:v for k,v in _r.items() if k not in ['svg','rows','next_token_rows','emotion_readouts','history']}
        _items=[mo.md("### Latest completed experiment: "+latest['experiment']),mo.md("Results belong to the submitted settings below. Edits are included only after you submit again."),mo.accordion({'Submitted settings':mo.md('```json\n'+json.dumps(latest['settings'],indent=2)+'\n```')})]
        if 'svg' in _r: _items.append(mo.Html(_r['svg']))
        if 'before' in _r:
            _before=_r['before'];_after=_r['after']
            _items.append(mo.md(f"**Before feedback:** {_before['wrong_answers']} wrong answers out of {_before['answered']} answers, from {_before['questions']} questions. **After feedback:** {_after['wrong_answers']} wrong answers out of {_after['answered']} answers. The policy asked to check on {_after['questions']-_after['answered']} questions. Compare both mistakes and willingness to answer."))
        if 'train_correct' in _r:
            _items.append(mo.md(f"**Training pairs:** {_r['train_correct']} / 80 correct. **Reserved pairs:** {_r['test_correct']} / 20 correct. Did learning the examples teach a rule that works on new pairs?"))
        if 'heldout_probe_rmse' in _r:
            _items.append(mo.md(f"**Reserved probe error:** {_r['heldout_probe_rmse']:.3f}. **Constant baseline error:** {_r['constant_baseline_rmse']:.3f}. Lower is better. A probe that loses to the baseline has not demonstrated reliable number information on these prompts."))
        if 'pca_variance' in _r:
            _items.append(mo.md(f"The two-dimensional map keeps **{100*_r['pca_variance']:.1f}%** of variation in the original reference vectors. Read the numbered rows below to connect points to evidence."))
        _items.append(mo.accordion({'Model, method and full measurements':mo.md('```json\n'+json.dumps(_overview,indent=2)+'\n```')}))
        _items.append(mo.ui.table(_r.get('rows',[]),page_size=12,label="Evidence table"))
        if _r.get('next_token_rows'): _items.append(mo.ui.table(_r['next_token_rows'],page_size=8,label="Actual next-token predictions"))
        if _r.get('emotion_readouts'): _items.extend([mo.md("**Separate prompted readout:** relative likelihoods among six candidate emotion labels, using mean token log likelihood. These are not the paper’s 0–10 ratings or calibrated emotional probabilities."),mo.ui.table(_r['emotion_readouts'])])
        _items.append(mo.download(data=json.dumps(latest,indent=2).encode(),filename="spark-evidence.json",label="Download settings and measured evidence"))
        display=mo.vstack(_items)
    display
    return (latest,)

@app.cell
def _(mo, lesson):
    reflection = mo.ui.text_area(label="After running: state your claim, cite two measurements, name a limitation, and propose a fair next test.",full_width=True)
    mo.vstack([mo.md("### Explain what you found\n"+lesson['assessment']),reflection,mo.md("**Teacher check:** look for a prediction, a controlled comparison, accurate use of measured evidence and a limit on the conclusion. Accept spoken or drawn explanations for younger students. Do not reward a desired result over an honest failed hypothesis.")])
    return (reflection,)

@app.cell
def _(mo, prediction, reflection, latest, json):
    lab_record = {"prediction":prediction.value,"reflection":reflection.value,"latest_experiment":latest}
    mo.download(data=json.dumps(lab_record,indent=2).encode(),filename="spark-lab-record.json",label="Save my prediction, evidence and explanation")
    return 

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