import { useState } from 'react'
import Plot from 'react-plotly.js'
import { Swords } from 'lucide-react'
import { runArena } from '../lib/api.js'
import { RewardBreakdown } from '../components/RewardBreakdown.jsx'

const SAMPLE = `class Optimizer:
    def __init__(self, dim):
        self.lr = 0.05
        self.beta = 0.9
        self.v = np.zeros(dim)

    def step(self, x, f_val, grad):
        # SGD with heavy-ball momentum
        self.v = self.beta * self.v - self.lr * grad
        return x + self.v
`

const TEMPLATES = [
  'quadratic', 'rosenbrock', 'styblinski_tang', 'huber',
  'gaussian_mix', 'himmelblau', 'plateau', 'cliff',
]

export function OptimizerArena() {
  const [template, setTemplate] = useState('quadratic')
  const [dim, setDim] = useState(5)
  const [seed, setSeed] = useState(42)
  const [code, setCode] = useState(SAMPLE)
  const [data, setData] = useState(null)
  const [loading, setLoading] = useState(false)
  const [err, setErr] = useState(null)

  async function run() {
    setLoading(true); setErr(null); setData(null)
    try {
      const r = await runArena({ template, dim, seed, code })
      if (r.error) {
        setErr(r.error)
      } else {
        setData(r)
      }
    } catch (e) {
      setErr(e.message)
    } finally {
      setLoading(false)
    }
  }

  return (
    <div className="grid grid-cols-[340px_1fr] gap-5">
      <aside className="card space-y-4 h-fit">
        <div>
          <h3 className="mb-1">Optimizer arena</h3>
          <p className="text-xs text-muted">
            Paste an <code>Optimizer</code> class. We run it through the full
            Phase-D arena vs tuned Adam on the chosen landscape.
            <br /><span className="text-subtle">np is pre-injected — no import lines.</span>
          </p>
        </div>
        <div>
          <label className="label">Template</label>
          <select className="input" value={template}
                  onChange={e => setTemplate(e.target.value)}>
            {TEMPLATES.map(t => <option key={t}>{t}</option>)}
          </select>
        </div>
        <div>
          <label className="label">Dim: {dim}</label>
          <input type="range" min={2} max={10} step={1} value={dim}
                 onChange={e => setDim(Number(e.target.value))}
                 className="w-full accent-accent" />
        </div>
        <div>
          <label className="label">Seed: {seed}</label>
          <input type="range" min={0} max={100} step={1} value={seed}
                 onChange={e => setSeed(Number(e.target.value))}
                 className="w-full accent-accent" />
        </div>
        <button className="btn-primary w-full" onClick={run} disabled={loading}>
          <Swords size={14} /> {loading ? 'Running…' : 'Run arena'}
        </button>
      </aside>

      <div className="space-y-4">
        <div className="card">
          <h3 className="mb-2">Your Optimizer class</h3>
          <textarea className="input font-mono text-xs h-72 leading-relaxed"
                    value={code} onChange={e => setCode(e.target.value)} spellCheck={false} />
        </div>
        {err && (
          <div className="card border-bad/40 bg-bad/10 text-bad">
            <strong>Compile error:</strong>
            <pre className="mt-1 text-xs whitespace-pre-wrap">{err}</pre>
          </div>
        )}
        {data?.summary_md && (
          <div className="card prose prose-invert max-w-none text-sm"
               dangerouslySetInnerHTML={{ __html: data.summary_md }} />
        )}
        {data && (
          <div className="grid grid-cols-2 gap-4">
            {data.contour && (
              <div className="card"><Plot data={data.contour.data} layout={data.contour.layout}
                    config={{ displayModeBar: false }} style={{ width: '100%' }} useResizeHandler /></div>
            )}
            {data.progress && (
              <div className="card"><Plot data={data.progress.data} layout={data.progress.layout}
                    config={{ displayModeBar: false }} style={{ width: '100%' }} useResizeHandler /></div>
            )}
          </div>
        )}
        {data?.breakdown && (
          <div className="card">
            <RewardBreakdown breakdown={data.breakdown} total={data.total} />
          </div>
        )}
      </div>
    </div>
  )
}