website/src/pages/DetectabilityPrediction.tsx

import { useState } from 'react'
import Card from '../components/Card'
import Button from '../components/Button'
import SequenceInput from '../components/SequenceInput'
import ModelPicker from '../components/ModelPicker'
import ConditionMatcher from '../components/ConditionMatcher'
import ResultsTable, { type Column } from '../components/ResultsTable'
import { findModel } from '../data/models'
import { classLabel } from '../data/screenMeta'
import type { Conditions } from '../lib/rankModels'
import { predictDetectability, screenBioactivityPeptides } from '../lib/api'
import { SAMPLE_PEPTIDES, type DetectabilityRow } from '../data/mockData'

const COLUMNS: Column<DetectabilityRow>[] = [
  {
    key: 'peptide',
    header: 'Peptide',
    render: (r) => <span className="mono">{r.peptide}</span>,
  },
  { key: 'length', header: 'Length', sortable: true, align: 'right' },
  {
    key: 'score',
    header: 'Detectability score',
    sortable: true,
    align: 'right',
    render: (r) => r.score.toFixed(3),
  },
  {
    key: 'prediction',
    header: 'Prediction',
    sortable: true,
    render: (r) =>
      r.prediction === 'Detectable' ? (
        <span className="badge badge-good">Detectable</span>
      ) : (
        <span className="badge badge-muted">Not detectable</span>
      ),
  },
  {
    key: 'bioactivities',
    header: 'Known bioactivity',
    render: (r) => {
      const list = (r.bioactivities ?? '').split(';').filter(Boolean)
      return list.length ? (
        <span className="badge badge-blue" title={list.map(classLabel).join('; ')}>
          {classLabel(list[0])}
          {list.length > 1 ? ` +${list.length - 1}` : ''}
        </span>
      ) : (
        <span className="muted">n/a</span>
      )
    },
  },
  {
    key: 'inTraining',
    header: 'In training',
    sortable: true,
    render: (r) =>
      r.inTraining === 'positive' ? (
        <span className="badge badge-warn" title="Detected (positive) example in this model's training set. The score may reflect memorisation.">
          Train⁺
        </span>
      ) : r.inTraining === 'negative' ? (
        <span className="badge badge-warn" title="Undetected (negative) example in this model's training set. The score may reflect memorisation.">
          Train⁻
        </span>
      ) : (
        <span className="muted">held out</span>
      ),
  },
  { key: 'model', header: 'Model', sortable: true },
]

export default function DetectabilityPrediction() {
  const [input, setInput] = useState('')
  const [modelCode, setModelCode] = useState('H9')
  const [conditions, setConditions] = useState<Conditions>({})
  const [pickMode, setPickMode] = useState<'match' | 'custom'>('match')
  const [rows, setRows] = useState<DetectabilityRow[] | null>(null)
  const [loading, setLoading] = useState(false)
  const [live, setLive] = useState(false)

  const canRun = input.trim().length > 0 && !loading
  const model = findModel(modelCode)

  async function run() {
    setLoading(true)
    try {
      // Predict detectability and screen known bioactivity in parallel, then join
      // by peptide so the output (and CSV) carries function alongside detectability.
      const [pred, bio] = await Promise.all([
        predictDetectability(input, model),
        screenBioactivityPeptides(input),
      ])
      const bioMap = new Map(
        bio.rows.map((b) => [b.peptide.toUpperCase(), b.bioactivities]),
      )
      setRows(
        pred.rows.map((r) => ({
          ...r,
          bioactivities: bioMap.get(r.peptide.toUpperCase()) ?? '',
        })),
      )
      setLive(pred.live)
    } finally {
      setLoading(false)
    }
  }

  const detectable = rows?.filter((r) => r.prediction === 'Detectable').length ?? 0

  return (
    <main className="page">
      <div className="page-header">
        <h1>Detectability Prediction</h1>
        <p className="lead">
          Paste peptides and pick a model. Describe your conditions to auto-rank the
          best fit, or choose one yourself. Peptides seen during training are flagged, and
          each peptide is also screened for known bioactivity.
        </p>
      </div>

      <div className="grid-2">
        <Card title="Peptides">
          <SequenceInput
            label="Peptide sequences"
            value={input}
            onChange={setInput}
            onLoadExample={() => setInput(SAMPLE_PEPTIDES)}
            placeholder={'PEPTIDE\nPEPTIDE\n\nor comma-separated / FASTA'}
          />
        </Card>
        <Card title="Model">
          <div className="segmented" style={{ marginBottom: 14, width: '100%' }}>
            <button
              className={pickMode === 'match' ? 'active' : ''}
              aria-pressed={pickMode === 'match'}
              onClick={() => setPickMode('match')}
              style={{ flex: 1 }}
            >
              Match my conditions
            </button>
            <button
              className={pickMode === 'custom' ? 'active' : ''}
              aria-pressed={pickMode === 'custom'}
              onClick={() => setPickMode('custom')}
              style={{ flex: 1 }}
            >
              Custom select a model
            </button>
          </div>
          {pickMode === 'match' ? (
            <ConditionMatcher
              conditions={conditions}
              onChange={setConditions}
              onPick={setModelCode}
              selected={modelCode}
            />
          ) : (
            <ModelPicker selected={modelCode} onChange={setModelCode} />
          )}
          <div className="selected-model">
            Selected model: <span className="mono">{model.code}</span> · {model.species.split('(')[0].trim()} · {model.enzyme}
          </div>
          <div style={{ marginTop: 12 }}>
            <Button onClick={run} disabled={!canRun}>
              {loading ? 'Predicting…' : 'Predict detectability'}
            </Button>
            {!canRun && (
              <p
                className="muted"
                style={{ fontSize: 'var(--text-sm)', marginTop: 8, marginBottom: 0 }}
              >
                Add at least one peptide.
              </p>
            )}
          </div>
        </Card>
      </div>

      <Card
        title="Detectability predictions"
        actions={
          rows != null && rows.length > 0 ? (
            <span
              className={`badge ${live ? 'badge-good' : 'badge-muted'}`}
              title={
                live
                  ? 'Scores computed by the model (server-side inference).'
                  : 'Backend unavailable. Showing illustrative offline scores.'
              }
            >
              {live ? 'Live inference' : 'Offline preview'}
            </span>
          ) : undefined
        }
      >
        {loading ? (
          <div className="empty-state">
            Scoring peptides… computing features and running the model.
          </div>
        ) : rows == null ? (
          <div className="empty-state">
            Run a prediction to see per-peptide detectability scores.
          </div>
        ) : rows.length === 0 ? (
          <div className="empty-state">No peptides parsed from the input.</div>
        ) : (
          <>
            <div className="stat-row">
              <div className="stat">
                <div className="stat-value">{rows.length}</div>
                <div className="stat-label">Peptides scored</div>
              </div>
              <div className="stat">
                <div className="stat-value">{detectable}</div>
                <div className="stat-label">Predicted detectable</div>
              </div>
              <div className="stat">
                <div className="stat-value" style={{ fontSize: '1.05rem' }}>
                  {model.code} · {model.enzyme}
                </div>
                <div className="stat-label">Model used ({model.species})</div>
              </div>
            </div>
            <ResultsTable
              rows={rows}
              columns={COLUMNS}
              csvFilename="hydropd_detectability.csv"
            />
          </>
        )}
      </Card>
    </main>
  )
}