crdc-demo/src/components/LoadingAnimation.jsx

import { useState, useEffect } from 'react'
import ChartPanel from './ChartPanel.jsx'
import * as api from '../hooks/useApi.js'
import { STUDENT_GROUPS, MODEL_QUADRANTS } from '../hooks/useApi.js'
import { shortGroupLabel } from '../utils/colors.js'

// CRDC waves to fetch (for Chart 1 — time series)
const WAVES = ['21-22', '17-18', '15-16']

/**
 * Loading screen with animated histogram grid.
 * Fetches all required data in parallel, shows progress as bars fill.
 * When complete → renders ChartPanel directly (no transition needed).
 */
export default function LoadingAnimation({ district, state }) {
  const [loadedCount, setLoadedCount] = useState(0)
  const [totalCalls, setTotalCalls] = useState(0)
  const [error, setError] = useState(null)

  // Build the full list of API calls needed for all 3 charts
  useEffect(() => {
    let cancelled = false

    async function loadData() {
      try {
        // ——— Chart 1: Arrests over time (3 waves × default model) ———
        const wavePromises = WAVES.map((year) =>
          fetchDistrictEstimatesBatch(district.leaid, year)
            .then(() => !cancelled && setLoadedCount(c => c + 1))
            .catch(() => { /* individual failure doesn't block */ })
        )

        // ——— Chart 2: Rate by group (8 groups × default model, most recent wave) ———
        const groupPromises = STUDENT_GROUPS.map((sg) =>
          fetchDistrictEstimatesBatch(district.leaid, '21-22', sg.race, sg.sex)
            .then(() => !cancelled && setLoadedCount(c => c + 1))
            .catch(() => {})
        )

        // ——— Chart 3: posterior density, all 4 quadrant models so the dropdown can switch (8 groups × 4 = 32 calls) ———
        const modelPromises = MODEL_QUADRANTS.flatMap((quad) =>
          STUDENT_GROUPS.map((sg) =>
            fetchDistrictEstimatesBatch(district.leaid, '21-22', sg.race, sg.sex, quad.model)
              .then(() => !cancelled && setLoadedCount(c => c + 1))
              .catch(() => {})
          )
        )

        // Set total before starting (for progress bar)
        const allPromises = [...wavePromises, ...groupPromises, ...modelPromises]
        if (!cancelled) setTotalCalls(allPromises.length)

        await Promise.all(allPromises)

        if (!cancelled) {
          // All data loaded — ChartPanel renders in place of this component
          // We use a render prop pattern: return <ChartPanel /> when done
          setLoadedCount(allPromises.length)
        }
      } catch (err) {
        if (!cancelled) setError(err.message || 'Failed to load data')
      }
    }

    loadData()
    return () => { cancelled = true }
  }, [district.leaid])

  // ——— Render the histogram grid animation or ChartPanel when done ———
  const isComplete = loadedCount >= totalCalls && totalCalls > 0

  if (isComplete) {
    return <ChartPanel district={district} state={state} />
  }

  if (error) {
    return renderError(error, district)
  }

  // Build the grid of "bars" — one per API call needed.
  // Total = WAVES (3) + STUDENT_GROUPS (8) + MODEL_QUADRANTS×STUDENT_GROUPS (4×8=32)
  const TOTAL_BARS = WAVES.length + STUDENT_GROUPS.length + (MODEL_QUADRANTS.length * STUDENT_GROUPS.length)

  // Build bar metadata for rendering
  const allBars = Array.from({ length: TOTAL_BARS }, (_, i) => {
    if (i < WAVES.length) return { subLabel: 'Time series', group: 'waves' }
    if (i < WAVES.length + STUDENT_GROUPS.length) {
      const sg = STUDENT_GROUPS[i - WAVES.length]
      return { subLabel: shortGroupLabel(sg.race, sg.sex), group: 'groups' }
    }
    const quadIdx = Math.floor((i - WAVES.length - STUDENT_GROUPS.length) / STUDENT_GROUPS.length)
    const sgIdx = (i - WAVES.length - STUDENT_GROUPS.length) % STUDENT_GROUPS.length
    const sg = STUDENT_GROUPS[sgIdx]
    return { subLabel: shortGroupLabel(sg.race, sg.sex), group: 'models', modelIdx: quadIdx }
  })

  return (
    <div style={{ padding: 'var(--space-5) 0 var(--space-7)' }}>
      {/* Header */}
      <span className="eyebrow">Fetching data from the CRDC Arrest Rate API</span>
      <h2 style={{ marginTop: 'var(--space-2)', marginBottom: 'var(--space-3)' }}>
        Building estimates for {district.lea_name} in {state === 'DC' ? "District of Columbia" : getStateName(state)}
      </h2>

      {/* Progress counter */}
      <div style={{ display: 'flex', justifyContent: 'space-between', marginBottom: 'var(--space-3)' }}>
        <span style={{ color: 'var(--cv-navy-600)', fontWeight: 600, fontSize: '1.1rem' }}>
          {loadedCount} of {totalCalls || allBars.length} datasets loaded
        </span>
        <span className="stat-label">
          {Math.round((loadedCount / (totalCalls || allBars.length)) * 100)}% complete
        </span>
      </div>

      {/* Animated histogram grid — mirrors the Bayesian posterior draw theme */}
      <div style={{
        display: 'grid',
        gridTemplateColumns: `repeat(${3}, 1fr)`,
        gap: 'var(--space-2)',
        marginTop: 'var(--space-4)'
      }}>
        {/* Column headers */}
        {['Time series (Chart 1)', 'By student group (Chart 2)', 'Posterior density (Chart 3)'].map((h, i) => (
          <div key={i} style={{ textAlign: 'center', paddingBottom: 'var(--space-1)' }}>
            <span className="stat-label" style={{ display: 'block' }}>{h}</span>
          </div>
        ))}

        {/* Grid of animated bars — each represents one API call */}
        {allBars.map((bar, i) => {
          const filled = i < loadedCount
          // Random height for "data viz" aesthetic (deterministic via seed = index)
          const heightSeed = ((i * 37) % 100) + 20  // 20–120px range
          return (
            <div key={i} style={{ display: 'flex', flexDirection: 'column', alignItems: 'center' }}>
              {/* The animated bar */}
              <div style={{
                width: '100%', height: `${heightSeed}px`,
                background: filled
                  ? `var(--cv-navy-${600 - (i % 2) * 100})`
                  : 'var(--paper-3)',
                border: '1px solid var(--cv-rule)',
                borderRadius: '4px 4px 0 0',
                transition: 'height 300ms ease, background 300ms ease',
                position: 'relative',
                overflow: 'hidden'
              }}>
                {filled && (
                  <div style={{
                    position: 'absolute', bottom: 0, left: 0, right: 0, height: '100%',
                    background: `linear-gradient(to top, var(--teal-600) 0%, transparent ${70 + (i % 30)}%)`,
                    opacity: 0.4, transition: 'opacity 500ms ease'
                  }} />
                )}
              </div>
              {/* Label below */}
              <span style={{ fontSize: '0.65rem', color: filled ? 'var(--cv-ink)' : 'var(--cv-ink-4)', marginTop: '0.25rem', textAlign: 'center' }}>
                {bar.subLabel}
              </span>
            </div>
          )
        })}
      </div>

      {/* Context message */}
      <p style={{ fontSize: '0.9rem', color: 'var(--cv-ink-2)', marginTop: 'var(--space-4)' }}>
        Fetching {totalCalls || allBars.length} datasets across 3 CRDC waves, 8 student groups, and 4 Bayesian model specifications...
      </p>

      {/* Model legend */}
      <div style={{ display: 'flex', gap: 'var(--space-2)', marginTop: 'var(--space-3)', flexWrap: 'wrap' }}>
        {MODEL_QUADRANTS.map((q, i) => (
          <span key={i} className="stat-label" style={{ display: 'inline-flex', alignItems: 'center', gap: '0.4rem' }}>
            <span style={{ width: 12, height: 12, background: `var(--cv-navy-${600 - i * 150})`, borderRadius: 2 }} />
            {q.label}
          </span>
        ))}
      </div>
    </div>
  )
}

// ——— API helper using the centralized client from useApi.js ———
// This ensures consistent error handling, envelope unwrapping, and CORS proxy support.
async function fetchDistrictEstimatesBatch(leaid, year, race = null, sex = null, model = 'unified_m2_mod') {
  return api.fetchDistrictEstimates(leaid, { year, model, race, sex })
}

function getStateName(code) {
  const names = {
    AL:'Alabama',AK:'Alaska',AZ:'Arizona',AR:'Arkansas',CA:'California',CO:'Colorado',CT:'Connecticut',DE:'Delaware',DC:'District of Columbia',FL:'Florida',GA:'Georgia',HI:'Hawaii',ID:'Idaho',IL:'Illinois',IN:'Indiana',IA:'Iowa',KS:'Kansas',KY:'Kentucky',LA:'Louisiana',ME:'Maine',MD:'Maryland',MA:'Massachusetts',MI:'Michigan',MN:'Minnesota',MS:'Mississippi',MO:'Missouri',MT:'Montana',NE:'Nebraska',NV:'Nevada',NH:'New Hampshire',NJ:'New Jersey',NM:'New Mexico',NY:'New York',NC:'North Carolina',ND:'North Dakota',OH:'Ohio',OK:'Oklahoma',OR:'Oregon',PA:'Pennsylvania',RI:'Rhode Island',SC:'South Carolina',SD:'South Dakota',TN:'Tennessee',TX:'Texas',UT:'Utah',VT:'Vermont',VA:'Virginia',WA:'Washington',WV:'West Virginia',WI:'Wisconsin',WY:'Wyoming'
  }
  return names[code] || code
}

function renderError(error, district) {
  return (
    <div style={{ padding: 'var(--space-5) 0 var(--space-7)', textAlign: 'center' }}>
      <span className="stat-label" style={{ color: 'var(--cv-danger)' }}>ERROR</span>
      <h3 style={{ marginTop: 'var(--space-2)' }}>Failed to load data for {district.lea_name}</h3>
      <p style={{ color: 'var(--cv-ink-2)', maxWidth: '40rem', margin: '0 auto var(--space-3)' }}>
        {error}. This may be a temporary API issue or the district may not have available estimates.
      </p>
      <button className="btn-outline" onClick={() => window.location.href = '/'}>
        ← Try another district
      </button>
    </div>
  )
}