crdc-demo/src/charts/ObservedRateDensity.jsx

import { useState } from 'react'

/**
 * Chart 5: Observed rate per group against model distribution.
 * Shows observed arrest rate and overlays the posterior draw distributions from all four models
 * as ridgeline-style density bars (proxy using interval bounds + median).
 * Mirrors wp_fig_group_density() / wp_fig_hpd_ridges() from paper_figures.R.
 */

const MODEL_LABELS = {
  unified_m1_mod: 'One-year, baseline',
  unified_m2_mod: 'One-year + covariate',
  unified_m3_mod: 'Three-year, baseline',
  unified_m4_mod: 'Three-year + covariate',
}

export default function ObservedRateDensity({ quadData }) {
  const [expandedGroup, setExpandedGroup] = useState(0) // first group by default

  if (!quadData || Object.keys(quadData).length === 0) {
    return (
      <div className="cv-card" style={{ padding: 'var(--space-2)' }}>
        <h3 style={{ fontSize: '0.85rem', marginBottom: 'var(--space-1)', color: 'var(--cv-ink-2)' }}>
          Observed rate vs. model distribution — by student group
        </h3>
        <p style={{ color: 'var(--cv-ink-3)' }}>No model data available.</p>
      </div>
    )
  }

  // Extract unique groups from the default model (unified_m2_mod)
  const defaultModel = quadData['unified_m2_mod'] || []
  const groups = [...new Map(defaultModel.map(r => [`${r.race}-${r.sex}`, { race: r.race, sex: r.sex, label: getLabel(r.race, r.sex) }])).values()]

  // Get data for the expanded group across all 4 models
  const currentGroup = groups[expandedGroup] || groups[0]
  const modelResults = Object.entries(quadData).map(([modelId, rows]) => {
    const row = rows.find(r => r.race === currentGroup?.race && r.sex === currentGroup?.sex)
    if (!row) return null

    // Convert rate median/interval to per-1k scale for visualization
    const enroll = row.stu_enroll || 0
    const observedRate = (row.observed_arrests || 0) / ((enroll || 1) / 1000)
    const modeledRateMedian = (row.rate_median || 0) * 1000   // rate is per student, convert to per 1k
    const modeledLower = (row.rate_lower || 0) * 1000
    const modeledUpper = (row.rate_upper || 0) * 1000

    return { modelId, observedRate, modeledRateMedian, modeledLower, modeledUpper }
  }).filter(Boolean)

  // Max value for scaling the ridgeline chart
  const maxRate = Math.max(...modelResults.map(r => r.modeledUpper), ...modelResults.map(r => r.observedRate), 0.5)

  const width = 280
  const height = 220
  const margin = { top: 30, right: 20, bottom: 60, left: 45 }

  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)' }}>
        Observed rate vs. model distribution — {currentGroup?.label || 'Select a group'}
      </h3>

      {/* Group selector */}
      <select value={expandedGroup} onChange={(e) => setExpandedGroup(Number(e.target.value))}
              style={{ marginBottom: 'var(--space-1)', padding: '0.25rem 0.5rem', fontFamily: 'var(--font-sans)' }}>
        {groups.map((g, i) => (
          <option key={`${g.race}-${g.sex}`} value={i}>{g.label}</option>
        ))}
      </select>

      {/* Ridgeline-style visualization (proxy using interval bars) */}
      <svg width="100%" viewBox={`0 0 ${width} ${height + modelResults.length * 25}`} style={{ maxWidth: '100%' }}>
        {modelResults.map((r, i) => {
          const rowY = margin.top + i * (height / modelResults.length + 8)
          const barHeight = height / modelResults.length - 6
          // Pre-compute coordinates outside JSX to avoid embedding statements in render output
          const obsWidth = (r.observedRate / maxRate) * (width - margin.left - margin.right)
          const lowerX = margin.left + (r.modeledLower / maxRate) * (width - margin.left - margin.right)
          const upperWidth = ((r.modeledUpper - r.modeledLower) / maxRate) * (width - margin.left - margin.right)
          const medX = margin.left + (r.modeledRateMedian / maxRate) * (width - margin.left - margin.right)

          return (
            <g key={r.modelId}>
              {/* Background band */}
              <rect x={margin.left} y={rowY} width={width - margin.left - margin.right} height={barHeight} rx={3} fill="var(--cv-paper-2)" />

              {/* Observed rate bar (solid) */}
              <rect x={margin.left} y={rowY + 4} width={obsWidth} height={barHeight - 8} fill="var(--cv-navy-600)" rx={2} opacity={0.9} />

              {/* Model interval (dashed outline showing 95% range) */}
              <rect x={lowerX} y={rowY + 4} width={upperWidth} height={barHeight - 8} fill="none" stroke="var(--teal-600)" strokeWidth={1.5} rx={2} />

              {/* Model median marker */}
              <line x1={medX} y1={rowY + 3} x2={medX} y2={rowY + barHeight - 3} stroke="var(--teal-600)" strokeWidth={2} />

              {/* Observed rate label */}
              {r.observedRate > maxRate * 0.05 && (
                <text x={margin.left + obsWidth + 4} y={rowY + barHeight / 2 + 3} fontSize="0.6rem" fill="var(--cv-ink)">{r.observedRate.toFixed(1)}</text>
              )}

              {/* Model label */}
              <text x={margin.left - 6} y={rowY + barHeight / 2 + 3} textAnchor="end" fontSize="0.65rem" fill="var(--cv-ink)">{MODEL_LABELS[r.modelId] || r.modelId}</text>

              {/* Y-axis gridline */}
              <line x1={margin.left - 4} y1={rowY + barHeight / 2 - 8} x2={margin.left} y2={rowY + barHeight / 2 - 8} stroke="var(--cv-rule)" strokeWidth={0.5} />
            </g>
          )
        })}

        {/* X-axis */}
        <line x1={margin.left} y1={height + modelResults.length * 16} x2={width - margin.right} y2={height + modelResults.length * 16} stroke="var(--cv-rule)" strokeWidth={1} />
        {[0, 0.5, 1].map((frac) => {
          const val = frac * maxRate
          return (
            <text key={`x-${val}`} x={margin.left + frac * (width - margin.left - margin.right)} y={height + modelResults.length * 24} textAnchor="middle" fontSize="0.6rem" fill="var(--cv-ink-3)">{val.toFixed(1)}</text>
          )
        })}
        <text x={margin.left + (width - margin.left - margin.right) / 2} y={height + modelResults.length * 30} textAnchor="middle" fontSize="0.7rem" fill="var(--cv-ink-3)">Rate per 1,000 students</text>

        {/* Legend */}
        <g transform={`translate(${margin.left}, ${height + modelResults.length * 8})`} style={{ fontSize: '0.65rem' }}>
          <rect x={0} y={0} width={12} height={12} fill="var(--cv-navy-600)" rx={2} />
          <text x={18} y={10}>Observed rate</text>
          <line x1={140} y1={6} x2={156} y2={6} stroke="var(--teal-600)" strokeWidth={2} />
          <circle cx={148} cy={6} r={3} fill="var(--teal-600)" />
          <text x={162} y={10}>Model (median + 95% interval)</text>
        </g>
      </svg>

      {/* Note */}
      <p style={{ fontSize: '0.72rem', color: 'var(--cv-ink-3)', marginTop: 'var(--space-1)' }}>
        Solid blue = observed rate per 1,000 students. Teal outline + median line shows each model's posterior distribution (95% interval). Dashed lines indicate the range of plausible values according to each Bayesian specification.
      </p>
    </div>
  )
}

function getLabel(race, sex) {
  const r = race === 'WH' ? 'White' : race === 'BL' ? 'Black' : race === 'HI' ? 'Hispanic' : 'American Indian/Alaskan Native'
  const s = sex === 'F' ? 'Female' : 'Male'
  return `${r} ${s}`
}