import {useEffect,useRef,useState} from 'react';
import {useMotion} from './Motion';
import {DualScene} from './DualScene';
import {beats,beatIndex,phase,residuePreviewCount} from './dualStory';
import {useScrollStory} from './useScrollStory';
import type {Example} from './types';
import config from '../t5config_large.json';

function MiniBlock({protein=false,decoder=false}:{protein?:boolean;decoder?:boolean}) {
 return <div className={`reading-block ${protein?'is-protein':'is-text'}`}>
  <div><strong>{protein?'Protein':'Text'} {decoder?'Decoder':'Encoder'}</strong><span>×{decoder?config.num_decoder_layers:config.num_layers}</span></div>
  <p>{decoder?'Masked self-attention':'Self-attention'}</p>{decoder&&<p className="reading-context">Cross-attention ← projected context</p>}<p>Feed-forward</p>
 </div>;
}
function Vector({protein=false}:{protein?:boolean}) {return <span className={`reading-vector ${protein?'is-protein':'is-text'}`} aria-label="Illustrative hidden-state vector">{Array.from({length:18},(_,i)=><i key={i}/>)}</span>;}
function ReadingVisual({index,example}:{index:number;example:Example}) {
 if (index===0) return <div className="reading-visual"><span className="story-label">REFERENCE DESCRIPTION</span><blockquote>{example.text}</blockquote><div className="reading-down">↓ <span>Encode the description</span></div><MiniBlock/></div>;
 if (index===1) return <div className="reading-visual"><div className="reading-projection"><Vector/><span>→</span><div><strong>Linear Wₜ→ₚ</strong><small>{config.d_model} → {config.d_model}</small></div><span>→</span><Vector protein/></div><p className="reading-key">Hidden states become the decoder’s context.</p><div className="reading-down">↓ <span>Read through cross-attention</span></div><MiniBlock protein decoder/></div>;
 if (index===2) return <div className="reading-visual"><MiniBlock protein decoder/><div className="reading-down">↓ <span>Autoregressive output</span></div><div className="reading-tiles reading-residue-output is-protein">{example.generatedProtein.slice(0,residuePreviewCount).split('').map((s,i)=><span key={i}>{s}</span>)}</div><p className="reading-key">First {residuePreviewCount} of {example.generatedProtein.length} generated residues.</p></div>;
 if (index===3) return <div className="reading-visual"><span className="story-label">REFERENCE SEQUENCE · INDEPENDENT INPUT</span><div className="reading-tiles is-protein">{example.protein.slice(0,16).split('').map((s,i)=><span key={i}>{s}</span>)}</div><div className="reading-down">↓ <span>Encode the original protein</span></div><MiniBlock protein/><div className="reading-down">↓</div><div className="reading-projection"><Vector protein/><span>→</span><div><strong>Linear Wₚ→ₜ</strong><small>Separate learned projection</small></div><span>→</span><Vector/></div></div>;
 if (index===4) return <div className="reading-visual"><MiniBlock decoder/><div className="reading-down">↓ <span>Recorded text output</span></div><div className="reading-tiles reading-words is-text">{example.generatedText.replace(/([.!?])(?=[A-Z])/g,'$1 ').split(/\s+/).map((word,i)=><span key={i}>{word}</span>)}</div></div>;
 return <div className="reading-visual reading-joint"><div><span className="story-label">TEXT → PROTEIN</span><p><b className="is-text">Text Encoder</b><span>↓ Wₜ→ₚ</span><b className="is-protein">Protein Decoder</b></p></div><div><span className="story-label">PROTEIN → TEXT</span><p><b className="is-protein">Protein Encoder</b><span>↓ Wₚ→ₜ</span><b className="is-text">Text Decoder</b></p></div><div className="reading-shared"><p className="is-text">Text Encoder ↔ Text Decoder<br/><strong>Same text self-attention weights</strong></p><p className="is-protein">Protein Encoder ↔ Protein Decoder<br/><strong>Same protein self-attention weights</strong></p><small>Corresponding blocks 1–24 within each modality.</small></div><div className="reading-loss">L = ½ L(text → protein)<br/> + ½ L(protein → text)</div></div>;
}
export function DualLearning({example}:{example:Example}) {
 const {enabled,reduced,paused:pagePaused}=useMotion();
 const [small,setSmall]=useState(()=>matchMedia('(max-width: 1100px), (max-height: 640px)').matches);
 const [readAll,setReadAll]=useState(false),[paused,setPaused]=useState(false);
 const section=useRef<HTMLDivElement>(null),stage=useRef<HTMLDivElement>(null);
 const flow=small||reduced||readAll;
 const {progress,seek}=useScrollStory(section,stage,enabled&&!paused,flow);
 const active=beatIndex(progress);
 useEffect(()=>{
  const query=matchMedia('(max-width: 1100px), (max-height: 640px)');
  const update=()=>setSmall(query.matches);
  query.addEventListener('change',update);
  return()=>query.removeEventListener('change',update);
 },[]);
 return <section id="dual-learning" className="section dual-story" aria-labelledby="dual-story-title" data-story-mode={flow?'reading':'scroll'}>
  <header className="story-intro">
   <div><p className="eyebrow"><span>01</span> INSIDE OUROBOROS</p><h2 id="dual-story-title">One example.<br/><span>Two ways to learn.</span></h2></div>
   <div><p>Follow a function into a protein.<br/>Then follow a protein into language.</p><span className="story-intro-hint">{flow?'SIX STEPS THROUGH DUAL LEARNING':'SCROLL TO EXPLORE · SCROLL BACK TO REVISIT'} <span aria-hidden="true">↓</span></span></div>
  </header>
  {flow?<div className="story-reading">
   {!small&&!reduced&&<button className="text-button" onClick={()=>{setReadAll(false);document.getElementById('dual-learning')?.scrollIntoView({behavior:'instant'});}}>Return to scroll experience ↑</button>}
   {beats.map((beat,i)=><article key={beat.label} className={`reading-step tone-${beat.color}`} id={`dual-step-${i+1}`}><div><span className="story-label">{beat.tag}</span><h3>{beat.title}</h3><p>{beat.body}</p><span className="story-reading-note">{beat.note}</span></div><ReadingVisual index={i} example={example}/></article>)}
  </div>:<div className="story-scroll" ref={section}>
   <div className="story-stage" ref={stage} data-progress={progress.toFixed(6)} data-beat={active} data-paused={paused||pagePaused}>
    <div className="story-topline"><span>THE DUAL LEARNING FRAMEWORK</span><div><button onClick={()=>setPaused(v=>!v)} disabled={pagePaused} aria-pressed={paused||pagePaused}>{paused||pagePaused?'Resume scene':'Pause scene'}</button><button onClick={()=>{setReadAll(true);document.getElementById('dual-learning')?.scrollIntoView({behavior:'instant'});}}>Read all steps ↗</button></div></div>
    <div className="story-composition">
     <aside className="story-copy">
      <div className="story-beats">{beats.map((beat,i)=>{
       const fadeIn=i===0?1:phase(progress,beat.at,beat.at+.015);
       const fadeOut=i===beats.length-1?1:1-phase(progress,beats[i+1].at-.016,beats[i+1].at);
       return <article key={beat.label} className={`story-beat tone-${beat.color}`} aria-hidden={active!==i} style={{opacity:fadeIn*fadeOut,transform:`translateY(${(1-fadeIn)*12-(1-fadeOut)*12}px)`}}>
        <span className="story-label">{beat.tag}</span><h3>{beat.title}</h3><p>{beat.body}</p><div className="story-note"><i/>{beat.note}</div>
        {i===5&&<div className="sharing-inspector" role="group" aria-label="Inspect shared parameters"><span>{progress<.945?'Text encoder ↔ text decoder':progress<.986?'Protein encoder ↔ protein decoder':'Two modality-specific parameter links'}</span>{([['Text weights',.924],['Protein weights',.967],['Both',1]] as const).map(([label,p])=><button key={label} tabIndex={active===5?0:-1} aria-pressed={p===1?progress>=.986:p<.945?progress<.945:progress>=.945&&progress<.986} onClick={()=>seek(p)}>{label}</button>)}</div>}
       </article>;
      })}</div>
      <div className="story-sample"><span className="story-label">THE EXAMPLE</span><strong>Small heat shock protein</strong><span>HSP20 · {example.id} · large-v5</span></div>
     </aside>
     <div className="story-figure"><div className="story-figure-grid" aria-hidden="true"/><DualScene progress={progress} example={example}/><div className="story-figure-legend"><span><i/> LANGUAGE</span><span><i/> PROTEIN</span><span>SCHEMATIC · NOT TO SCALE</span></div></div>
    </div>
    <nav className="story-chapters" aria-label="Explanation chapters">{beats.map((beat,i)=><button key={beat.label} aria-current={active===i?'step':undefined} onClick={()=>seek(beat.sample)}><span>{String(i+1).padStart(2,'0')}</span>{beat.label}<i style={{transform:`scaleX(${phase(progress,beat.at,i===5?1:beats[i+1].at)})`}}/></button>)}</nav>
    <div className="story-scrubber"><span>{paused||pagePaused?'PAUSED · USE CHAPTERS TO STEP':'SCROLL TO ADVANCE OR DRAG'}</span><input type="range" min="0" max="1000" value={Math.round(progress*1000)} onChange={e=>seek(Number(e.target.value)/1000)} aria-label="Explanation progress" aria-valuetext={`Step ${active+1} of 6: ${beats[active].label}`}/><output>{String(active+1).padStart(2,'0')} / 06</output></div>
   </div>
  </div>}
  <div className="story-provenance"><span className="saved-badge">RECORDED EXAMPLE</span><p>Two independent predictions from the reference description and reference sequence. Animation timing is illustrative; training uses paired supervision, not cycle reconstruction.</p><a href="/data/core-example.json" download>Source data ↗</a></div>
  <details className="story-archive"><summary>Inspect the complete example <span>Reference inputs & archived predictions +</span></summary><div className="story-archive-grid"><article><span className="story-label">01 / TEXT → PROTEIN</span><h4>Reference description</h4><p>{example.text}</p><h4>Generated protein · {example.generatedProtein.length} residues</h4><p className="archive-sequence" data-archive="generated-protein">{example.generatedProtein}</p></article><article><span className="story-label">02 / PROTEIN → TEXT</span><h4>Reference protein · {example.protein.length} residues</h4><p className="archive-sequence" data-archive="reference-protein">{example.protein}</p><h4>Generated description</h4><p className="archive-text" data-archive="generated-text">{example.generatedText}</p></article></div></details>
  <details className="story-technical"><summary>Architecture & training notes <span>+</span></summary><p>Each large-v5 encoder and decoder contains {config.num_layers} T5-style blocks with hidden size {config.d_model.toLocaleString('en-US')}. Encoders use self-attention and feed-forward layers; decoders add masked self-attention and cross-attention to the projected source states. The diagram expands a representative block, omitting residual connections and layer normalization for clarity. Hidden-state squares illustrate vectors; their values are not measured activations.</p><p>The original model jointly optimizes two autoregressive, teacher-forced cross-entropy losses. Corresponding encoder/decoder self-attention parameters are shared within each modality. The two linear projections are initialized as transposes and optimized independently. The displayed predictions are archived inference outputs, not a visualization of teacher forcing. Word chips are readable excerpts, not T5 subword boundaries.</p></details>
 </section>;
}