import { useMemo, useState } from 'react'
import { MODELS, PFLY_OOB, displayAuroc, featureSetLabel, type RealModel } from '../data/models'
import ModelDetail from './ModelDetail'

interface ModelPickerProps {
  selected: string // model code
  onChange: (code: string) => void
}

interface Group {
  key: string
  species: string
  models: RealModel[]
}

export default function ModelPicker({ selected, onChange }: ModelPickerProps) {
  const [openInfo, setOpenInfo] = useState<string | null>(null)

  // group models by species
  const groups = useMemo<Group[]>(() => {
    const byKey = new Map<string, Group>()
    for (const m of MODELS) {
      if (!byKey.has(m.speciesKey)) byKey.set(m.speciesKey, { key: m.speciesKey, species: m.species, models: [] })
      byKey.get(m.speciesKey)!.models.push(m)
    }
    return [...byKey.values()]
  }, [])

  const selectedGroupKey = useMemo(
    () => MODELS.find((m) => m.code === selected)?.speciesKey ?? groups[0]?.key,
    [selected, groups],
  )

  // collapsed species sections, only the selected model's species is open initially
  const [openGroups, setOpenGroups] = useState<Set<string>>(new Set(selectedGroupKey ? [selectedGroupKey] : []))
  const toggleGroup = (k: string) =>
    setOpenGroups((prev) => {
      const next = new Set(prev)
      next.has(k) ? next.delete(k) : next.add(k)
      return next
    })

  return (
    <div className="field">
      <div className="field-label">
        Model{' '}
        <span className="field-hint">
          Pick Pfly (any species) or one of {MODELS.length} context-specific models grouped by species.
        </span>
      </div>
      <div className="option-list" style={{ marginBottom: 8 }}>
        <label className={'option' + (selected === PFLY_OOB.code ? ' selected' : '')}>
          <input
            type="radio"
            name="model"
            checked={selected === PFLY_OOB.code}
            onChange={() => onChange(PFLY_OOB.code)}
          />
          <span className="option-body">
            <span className="option-title" style={{ display: 'flex', gap: 8, alignItems: 'center' }}>
              <span className="mono">Pfly</span> · out-of-the-box
              <span className="badge badge-muted">sequence-only</span>
            </span>
            <span className="option-desc">
              Pretrained dlomix DetectabilityModel. Works for any species; not calibrated to a
              specific dataset.
            </span>
          </span>
        </label>
      </div>
      <div className="species-accordion">
        {groups.map((g) => {
          const isOpen = openGroups.has(g.key)
          const hasSelected = g.models.some((m) => m.code === selected)
          return (
            <div key={g.key} className={'species-group' + (isOpen ? ' open' : '')}>
              <button type="button" className="species-head" onClick={() => toggleGroup(g.key)}>
                <span className={'species-caret caret' + (isOpen ? ' open' : '')}>▶</span>
                <span className="species-name">{g.species.split('(')[0].trim()}</span>
                <span className="species-count">{g.models.length} model{g.models.length > 1 ? 's' : ''}</span>
                {hasSelected && !isOpen && <span className="badge badge-good species-sel">selected</span>}
              </button>
              {isOpen && (
                <div className="option-list species-body">
                  {g.models.map((m) => {
                    const on = selected === m.code
                    return (
                      <div key={m.code}>
                        <label className={'option' + (on ? ' selected' : '')}>
                          <input type="radio" name="model" checked={on} onChange={() => onChange(m.code)} />
                          <span className="option-body">
                            <span className="option-title" style={{ display: 'flex', gap: 8, alignItems: 'center' }}>
                              <span className="mono">{m.code}</span> · {m.enzyme}
                              <span className="badge badge-accent">AUROC {displayAuroc(m)}</span>
                            </span>
                            <span className="option-desc">
                              {featureSetLabel(m.bestFeatureSet)} · n={m.nPos.toLocaleString()}
                              {m.flags.length > 0 && ' · caveat, see details'}
                            </span>
                          </span>
                          <button
                            type="button"
                            className="btn btn-ghost btn-sm"
                            onClick={(e) => {
                              e.preventDefault()
                              setOpenInfo(openInfo === m.code ? null : m.code)
                            }}
                          >
                            {openInfo === m.code ? 'Hide' : 'Details'}
                          </button>
                        </label>
                        {openInfo === m.code && <ModelDetail m={m} />}
                      </div>
                    )
                  })}
                </div>
              )}
            </div>
          )
        })}
      </div>
    </div>
  )
}