website/src/pages/DetectabilityTraining.tsx

import { useState } from 'react'
import Card from '../components/Card'
import Button from '../components/Button'
import Field from '../components/Field'
import SequenceInput from '../components/SequenceInput'
import TaxIdInput from '../components/TaxIdInput'
import {
  CLASSIFIER_OPTIONS,
  DEFAULT_TRAINING_CONFIG,
  FEATURE_SET_OPTIONS,
  SPLIT_OPTIONS,
  SAMPLE_PEPTIDES,
  type TrainingConfig,
  type TrainingResult,
} from '../data/mockData'
import { trainDetectability } from '../lib/api'

export default function DetectabilityTraining() {
  const [input, setInput] = useState('')
  const [taxId, setTaxId] = useState('')
  const [useDefaults, setUseDefaults] = useState(true)
  const [cfg, setCfg] = useState<TrainingConfig>(DEFAULT_TRAINING_CONFIG)
  const [training, setTraining] = useState(false)
  const [result, setResult] = useState<TrainingResult | null>(null)
  const [live, setLive] = useState(false)
  const [downloadUrl, setDownloadUrl] = useState<string | null>(null)
  const [dataUrl, setDataUrl] = useState<string | null>(null)
  const [errorMsg, setErrorMsg] = useState<string | null>(null)

  const canRun = input.trim().length > 0 && taxId.length > 0 && !training

  function set<K extends keyof TrainingConfig>(k: K, v: TrainingConfig[K]) {
    setUseDefaults(false)
    setCfg((c) => ({ ...c, [k]: v }))
  }

  // Pfly is a coupled choice: picking it in either the classifier or the feature-set
  // dropdown keeps the other in sync (and de-selecting it restores the default).
  function setClassifier(v: string) {
    setUseDefaults(false)
    setCfg((c) => ({
      ...c,
      classifier: v,
      featureSet:
        v === 'pfly'
          ? 'pfly'
          : c.featureSet === 'pfly'
            ? DEFAULT_TRAINING_CONFIG.featureSet
            : c.featureSet,
    }))
  }
  function setFeatureSet(v: string) {
    setUseDefaults(false)
    setCfg((c) => ({
      ...c,
      featureSet: v,
      classifier:
        v === 'pfly'
          ? 'pfly'
          : c.classifier === 'pfly'
            ? DEFAULT_TRAINING_CONFIG.classifier
            : c.classifier,
    }))
  }

  function toggleDefaults(on: boolean) {
    setUseDefaults(on)
    if (on) setCfg(DEFAULT_TRAINING_CONFIG)
  }

  async function train() {
    setTraining(true)
    setResult(null)
    setErrorMsg(null)
    setDownloadUrl(null)
    setDataUrl(null)
    try {
      const out = await trainDetectability(input, taxId, cfg)
      setResult(out.result)
      setLive(out.live)
      setDownloadUrl(out.downloadUrl)
      setDataUrl(out.dataUrl)
      setErrorMsg(out.error ?? null)
    } finally {
      setTraining(false)
    }
  }

  return (
    <main className="page">
      <div className="page-header">
        <h1>Train a Detectability Model</h1>
        <p className="lead">
          Give detected peptides and a species tax ID, pick your specs or use the
          defaults, and train a model. Then download the model and processed data.
        </p>
      </div>

      <div className="grid-2">
        <Card title="Training data">
          <SequenceInput
            label="Detected peptide sequences (positives)"
            value={input}
            onChange={setInput}
            onLoadExample={() => setInput(SAMPLE_PEPTIDES)}
          />
          <TaxIdInput value={taxId} onChange={setTaxId} />
        </Card>

        <Card
          title="Model specifications"
          actions={
            <label
              className="option"
              style={{ padding: '6px 10px', marginBottom: 0 }}
            >
              <input
                type="checkbox"
                checked={useDefaults}
                onChange={(e) => toggleDefaults(e.target.checked)}
              />
              <span className="option-title" style={{ fontSize: 'var(--text-sm)' }}>
                Use defaults
              </span>
            </label>
          }
        >
          <Field label="Classifier">
            <select
              value={cfg.classifier}
              onChange={(e) => setClassifier(e.target.value)}
            >
              {CLASSIFIER_OPTIONS.map((o) => (
                <option key={o.id} value={o.id}>
                  {o.label} · {o.desc}
                </option>
              ))}
            </select>
          </Field>
          <Field
            label="Feature set"
            hint={
              cfg.featureSet === 'pfly'
                ? 'Pfly reads the sequence directly, no separate classifier'
                : undefined
            }
          >
            <select
              value={cfg.featureSet}
              onChange={(e) => setFeatureSet(e.target.value)}
            >
              {FEATURE_SET_OPTIONS.map((o) => (
                <option key={o.id} value={o.id}>
                  {o.label} · {o.desc}
                </option>
              ))}
            </select>
          </Field>
          <Field label="Split (cross-validation grouping)">
            <select
              value={cfg.split}
              onChange={(e) => set('split', e.target.value)}
            >
              {SPLIT_OPTIONS.map((o) => (
                <option key={o.id} value={o.id}>
                  {o.label} · {o.desc}
                </option>
              ))}
            </select>
          </Field>
          <div className="grid-2" style={{ gap: 14 }}>
            <Field label="CV folds">
              <input
                type="number"
                min={2}
                max={10}
                value={cfg.folds}
                onChange={(e) => set('folds', Number(e.target.value))}
              />
            </Field>
            <Field label="Negative balance">
              <select
                value={cfg.balance}
                onChange={(e) => set('balance', e.target.value)}
              >
                {['1:1', '1:3', '1:5'].map((b) => (
                  <option key={b} value={b}>
                    {b}
                  </option>
                ))}
              </select>
            </Field>
          </div>

          <div style={{ marginTop: 8 }}>
            <Button onClick={train} disabled={!canRun}>
              {training ? 'Training…' : 'Train model'}
            </Button>
            {!canRun && !training && (
              <p
                className="muted"
                style={{ fontSize: 'var(--text-sm)', marginTop: 8, marginBottom: 0 }}
              >
                Add positive peptides and a species tax ID to train.
              </p>
            )}
          </div>
        </Card>
      </div>

      <Card title="Training result">
        {training ? (
          <div className="empty-state">
            Training model… computing features, sampling k-mer negatives, and
            running cross-validation.
          </div>
        ) : result == null ? (
          <div className="empty-state">
            Train a model to see cross-validation metrics and per-fold results.
          </div>
        ) : (
          <>
            <p style={{ marginTop: 0, marginBottom: 12 }}>
              <span className={'badge ' + (live ? 'badge-good' : 'badge-muted')}>
                {live ? 'Live training' : 'Offline preview'}
              </span>{' '}
              <span className="muted" style={{ fontSize: 'var(--text-sm)' }}>
                {live
                  ? 'Trained on the server: assign proteins, sample negatives, protein-grouped CV.'
                  : 'Illustrative numbers; the training backend is not reachable.'}
              </span>
            </p>
            {errorMsg && (
              <p className="error-banner" style={{ marginTop: 0, marginBottom: 12 }}>
                Backend could not train on this input: {errorMsg}
              </p>
            )}
            <div className="stat-row">
              <div className="stat">
                <div className="stat-value">{result.rocAuc.toFixed(3)}</div>
                <div className="stat-label">
                  Mean ROC-AUC ({cfg.folds}-fold CV)
                </div>
              </div>
              <div className="stat">
                <div className="stat-value">{result.prAuc.toFixed(3)}</div>
                <div className="stat-label">Mean PR-AUC</div>
              </div>
              <div className="stat">
                <div className="stat-value">{result.nPositives}</div>
                <div className="stat-label">Positives</div>
              </div>
              <div className="stat">
                <div className="stat-value">{result.nNegatives}</div>
                <div className="stat-label">Negatives (sampled)</div>
              </div>
            </div>

            <div className="table-wrap">
              <table className="table">
                <thead>
                  <tr>
                    <th>Fold</th>
                    <th className="num">ROC-AUC</th>
                    <th className="num">PR-AUC</th>
                    <th className="num"># train</th>
                    <th className="num"># test</th>
                  </tr>
                </thead>
                <tbody>
                  {result.folds.map((f) => (
                    <tr key={f.fold}>
                      <td>Fold {f.fold}</td>
                      <td className="num">{f.rocAuc.toFixed(3)}</td>
                      <td className="num">{f.prAuc.toFixed(3)}</td>
                      <td className="num">{f.nTrain}</td>
                      <td className="num">{f.nTest}</td>
                    </tr>
                  ))}
                </tbody>
              </table>
            </div>

            <div style={{ marginTop: 16, display: 'flex', gap: 10, flexWrap: 'wrap' }}>
              {downloadUrl ? (
                <a href={downloadUrl} download>
                  <Button variant="secondary">Download trained model (.joblib)</Button>
                </a>
              ) : (
                <Button variant="secondary" disabled>
                  Download trained model (train on the backend to enable)
                </Button>
              )}
              {dataUrl && (
                <a href={dataUrl} download>
                  <Button variant="secondary">Download processed data (.csv)</Button>
                </a>
              )}
            </div>
            {dataUrl && (
              <p className="muted" style={{ fontSize: 'var(--text-xs)', marginTop: 8, marginBottom: 0 }}>
                The processed dataset has one row per peptide (your detected peptides + sampled
                negatives) with its assigned protein, bioactivity / cytotoxicity / allergen screening,
                and the full feature matrix (physicochemical + PepBERT + ESM2 embeddings).
              </p>
            )}
          </>
        )}
      </Card>
    </main>
  )
}