import type {CSSProperties} from 'react';
import modelConfig from '../t5config_large.json';

type Modality='text'|'protein';
export function ArchitectureNode({modality,kind,excerpt,active}:{modality:Modality;kind:'encoder'|'decoder';excerpt:string;active:boolean}){
 const decoder=kind==='decoder',protein=modality==='protein';
 const items=protein?excerpt.slice(0,16).split(''):excerpt.replace(/([.!?])(?=[A-Z])/g,'$1 ').split(/\s+/).slice(0,5);
 const layers=decoder?modelConfig.num_decoder_layers:modelConfig.num_layers;
 return <div className={`model-node node-${modality}-${kind}`} data-active={active}>
  <div className="architecture-node-title"><div><small>{decoder?'WRITE':'READ'}</small><strong>{protein?'Protein':'Text'} {kind}</strong></div><span className="stack-count">×{layers}</span></div>
  {!decoder&&<div className={`module-excerpt ${protein?'residue-excerpt':'text-excerpt'}`}><span className="excerpt-label">INPUT EXCERPT</span><div className="excerpt-tiles">{items.map((s,i)=><span className="input-tile" key={i}>{s}</span>)}<span className="excerpt-ellipsis">…</span></div></div>}
  <div className="transformer-stack" aria-label={`${layers} repeated transformer blocks`}>
   <div className="transformer-row attention-row"><span>{decoder?'Masked self-attention':'Self-attention'}</span><small>Q K V</small></div>
   {decoder&&<div className="transformer-row cross-attention-row"><span>Cross-attention</span><small>CONTEXT</small></div>}
   <div className="transformer-row feedforward-row"><span>Feed-forward</span><small>FFN</small></div>
  </div>
  {decoder&&<div className={`module-excerpt output-excerpt ${protein?'residue-excerpt':'text-excerpt'}`} data-emitting={active}><span className="excerpt-label">RECORDED OUTPUT EXCERPT</span><div className="excerpt-tiles">{items.map((s,i)=><span className="output-tile" key={i} style={{'--token-delay':`${i*.17}s`} as CSSProperties}>{s}</span>)}<span className="excerpt-ellipsis">…</span></div></div>}
 </div>;
}
export function Projection({reverse=false}:{reverse?:boolean}){
 return <div className={`projection-channel ${reverse?'projection-reverse':''}`} aria-label={reverse?'Protein-to-text linear projection':'Text-to-protein linear projection'}><div className="projection-line" aria-hidden="true"><i/><i/><i/></div><div className="projection-block"><span>Linear</span><strong>{reverse?'Wₚ→ₜ':'Wₜ→ₚ'}</strong></div><span className="projection-shape">{modelConfig.d_model} → {modelConfig.d_model}</span></div>;
}