import { useMemo } from 'react'

/**
 * Chart 4: Predicted draws by year vs. observed — for three-year models only.
 * Shows model's predicted median and 95% interval alongside the observed value,
 * faceted by four quadrants (1yr baseline / 1yr covariate / 3yr baseline / 3yr covariate).
 * Mirrors wp_fig_district_intervals() from paper_figures.R.
 */

const QUADRANT_LABELS = [
  { key: 'unified_m1_mod', label: 'One-year, no covariate' },
  { key: 'unified_m2_mod', label: 'One-year + referral rate' },
  { key: 'unified_m3_mod', label: 'Three-year, no covariate' },
  { key: 'unified_m4_mod', label: 'Three-year + referral rate' },
]

const WAVE_LABELS = { '15-16': '2015–16', '17-18': '2017–18', '21-22': '2021–22' }

export default function ModelDrawsComparison({ data, waves, districtName }) {
  // Derive per-wave observed arrest totals from the time series
  const waveData = useMemo(() => {
    return ['15-16', '17-18', '21-22'].map(year => {
      const rows = (data.waves && data.waves[year]) || []
      return {
        year, label: WAVE_LABELS[year],
        arrests: rows.reduce((sum, r) => sum + (r.observed_arrests || 0), 0),
        enroll: rows.reduce((sum, r) => sum + (r.stu_enroll || 0), 0),
      }
    })
  }, [data.waves])

  const quadModels = data.quadData || {}

  // Compute max arrests for Y-axis scaling across all quadrants and years
  const allArrestValues = []
  waveData.forEach(d => allArrestValues.push(d.arrests))
  Object.entries(quadModels).forEach(([modelId, rows]) => {
    ['15-16', '17-18', '21-22'].forEach(year => {
      const yearRow = (rows || []).find(r => r.year === year)
      if (yearRow) allArrestValues.push(Math.round((yearRow.count_median || 0)))
    })
  })
  const maxArrests = Math.max(...allArrestValues, 1)

  return (
    <div className="cv-card" style={{ padding: 'var(--space-2)' }}>
      {/* Chart title */}
      <h3 style={{ fontSize: '0.85rem', marginBottom: 'var(--space-1)', color: 'var(--cv-ink-2)' }}>
        Model predictions vs. observed — by wave & model type
      </h3>

      {/* Four small multiples (one per quadrant) */}
      <div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr', gap: 'var(--space-2)' }}>
        {QUADRANT_LABELS.map((q) => (
          <QuadrantChart key={q.key} q={q} waveData={waveData} quadModels={quadModels} maxArrests={maxArrests} />
        ))}
      </div>

      {/* Note */}
      <p style={{ fontSize: '0.72rem', color: 'var(--cv-ink-3)', marginTop: 'var(--space-1)' }}>
        Red circles = observed arrests. Teal points + error bars (three-year models only) show Bayesian model median and 90% interval. One-year models can only compare against the most recent wave.
      </p>
    </div>
  )
}

/** Single quadrant chart: predicted vs observed for one model specification */
function QuadrantChart({ q, waveData, quadModels, maxArrests }) {
  const width = 360
  const height = 280
  const margin = { top: 45, right: 25, bottom: 70, left: 55 }

  // For three-year models show all 3 waves; for one-year only the most recent
  const showYears = (q.key.includes('m3') || q.key.includes('m4')) ? ['15-16', '17-18', '21-22'] : ['21-22']

  // Derive chart dimensions from margin + width/height
  const innerWidth = width - margin.left - margin.right
  const innerHeight = height - margin.top - margin.bottom

  return (
    <div style={{ border: '1px solid var(--cv-rule)', borderRadius: 'var(--radius-md)', padding: 'var(--space-1)' }}>
      {/* Quadrant label */}
      <span className="stat-label" style={{ display: 'block', marginBottom: 'var(--space-0.5)' }}>{q.label}</span>

      <svg width="100%" viewBox={`0 0 ${width} ${height}`} style={{ maxWidth: '100%' }}>
        {/* Background */}
        <rect x={margin.left} y={margin.top} width={innerWidth} height={innerHeight} fill="var(--cv-paper-2)" rx={3} />

        {/* Y-axis gridlines + labels */}
        {[0, 0.25, 0.5, 0.75, 1].map((frac) => {
          const val = frac * maxArrests
          if (val > maxArrests) return null
          const y = margin.top + innerHeight - (val / Math.max(maxArrests, 1)) * innerHeight
          return (
            <g key={`qy-${frac}`}>
              <line x1={margin.left} y1={y} x2={margin.left + innerWidth} y2={y} stroke="var(--cv-rule)" strokeWidth={0.5} />
              {val > 0 && (
                <text x={margin.left - 6} y={y + 3} textAnchor="end" fontSize="0.6rem" fill="var(--cv-ink-4)">{Math.round(val)}</text>
              )}
            </g>
          )
        })}

        {/* Y-axis label */}
        <text x={12} y={margin.top + innerHeight / 2} textAnchor="middle" fontSize="0.6rem" fill="var(--cv-ink-4)" transform={`rotate(-90 12 ${margin.top + innerHeight / 2})`}>Arrests</text>

        {/* Points for each year */}
        {showYears.map((year, yi) => {
          const observed = waveData.find(d => d.year === year) || { arrests: 0 }
          const cx = margin.left + (yi / Math.max(showYears.length - 1, 1)) * innerWidth

          // Observed point position
          const observedY = margin.top + innerHeight - (observed.arrests / Math.max(maxArrests, 1)) * innerHeight

          // For three-year models: show predicted median + interval
          const isThreeYear = q.key.includes('m3') || q.key.includes('m4')
          if (!isThreeYear) {
            return (
              <g key={`yr-${year}-${q.key}`}>
                {/* Observed only for one-year models */}
                <circle cx={cx} cy={observedY} r={4} fill="var(--cv-danger)" stroke="#fff" strokeWidth={1} />
                <text x={cx} y={observedY - 10} textAnchor="middle" fontSize="0.6rem" fill="var(--cv-ink)">{Math.round(observed.arrests)}</text>
              </g>
            )
          }

          // Three-year model: show both observed (dodged left) and predicted (median + 95% interval, dodged right)
          const obsX = cx - 8
          const predX = cx + 8

          const yearRow = (quadModels[q.key] || []).find(r => r.year === year)
          if (!yearRow) {
            return <circle key={`obs-${year}`} cx={cx} cy={observedY} r={4} fill="var(--cv-danger)" />
          }

          // Compute predicted values scaled to the same enrollment as observed
          const enroll = yearRow.stu_enroll || 1
          const predMedian = (yearRow.count_median || 0) / (enroll / 1000) * (observed.enroll / 1000)
          const lower = (yearRow.count_lower || 0) / (enroll / 1000) * (observed.enroll / 1000)
          const upper = (yearRow.count_upper || 0) / (enroll / 1000) * (observed.enroll / 1000)

          return (
            <g key={`yr-${year}-${q.key}`}>
              {/* Observed */}
              <circle cx={obsX} cy={observedY} r={4} fill="var(--cv-danger)" stroke="#fff" strokeWidth={1} />
              {/* Predicted error bar + median point */}
              <line x1={predX} y1={margin.top + innerHeight - (lower / Math.max(maxArrests, 1)) * innerHeight}
                    x2={predX} y2={margin.top + innerHeight - (upper / Math.max(maxArrests, 1)) * innerHeight}
                    stroke="var(--teal-600)" strokeWidth={1.5} />
              <circle cx={predX} cy={margin.top + innerHeight - (predMedian / Math.max(maxArrests, 1)) * innerHeight} r={4} fill="var(--teal-600)" stroke="#fff" strokeWidth={1} />
            </g>
          )
        })}

        {/* X-axis labels */}
        {showYears.map((year, yi) => (
          <text key={`xl-${year}`} x={margin.left + (yi / Math.max(showYears.length - 1, 1)) * innerWidth} y={height - margin.bottom + 15} textAnchor="middle" fontSize="0.6rem" fill="var(--cv-ink)">{WAVE_LABELS[year]}</text>
        ))}

        {/* Legend */}
        <g transform={`translate(${margin.left}, ${height - 40})`}>
          <circle cx={0} cy={10} r={3} fill="var(--cv-danger)" />
          <text x={10} y={14} fontSize="0.6rem" fill="var(--cv-ink)">Observed</text>
          {q.key.includes('m3') && (
            <>
              <line x1={70} y1={7} x2={86} y2={7} stroke="var(--teal-600)" strokeWidth={1.5} />
              <circle cx={78} cy={7} r={3} fill="var(--teal-600)" />
              <text x={92} y={11} fontSize="0.6rem" fill="var(--cv-ink)">Model (median + 90%)</text>
            </>
          )}
        </g>
      </svg>
    </div>
  )
}