World population-density choropleth

import { defineChart } from '@tanstack/charts'
import { geoShape } from '@tanstack/charts/geo'
import { geoEqualEarth } from 'd3-geo'
import { scaleQuantize } from 'd3-scale'
import { worldLand, worldSphere } from '../../shared/fixtures/country-atlas'
import { learningPovertyCountries } from '../../shared/transforms/learning-poverty'
import { tanstackMount } from '../../shared/mount'
import type { ConformanceInput } from '../../types'

const colorRanges = [
  ['#ecfeff', '#a5f3fc', '#67e8f9', '#06b6d4', '#0e7490', '#164e63'],
  ['#f0fdf4', '#bbf7d0', '#86efac', '#22c55e', '#15803d', '#14532d'],
]
const projection = {
  type: geoEqualEarth,
  fit: 'sphere' as const,
}

const definition = (input: ConformanceInput) =>
  defineChart({
    marks: [
      geoShape([worldLand], {
        projection,
        fill: '#e2e8f0',
        stroke: '#ffffff',
        strokeWidth: 0.55,
      }),
      geoShape(learningPovertyCountries, {
        projection,
        color: (country) => country.properties.density,
        stroke: 'currentColor',
        strokeOpacity: 0.34,
        strokeWidth: 0.55,
      }),
      geoShape([worldSphere], {
        projection,
        fill: 'none',
        stroke: 'currentColor',
        strokeOpacity: 0.35,
        strokeWidth: 0.75,
      }),
    ],
    color: {
      scale: scaleQuantize<string>,
      range: colorRanges[input.revision % 2] ?? colorRanges[0],
    },
    margin: 12,
  })

export const mount = tanstackMount(
  definition,
  'World population-density choropleth',
  {
    format: ({ datum }) =>
      'properties' in datum && 'density' in datum.properties
        ? `${datum.properties['Country Name']} · ${datum.properties.density} people/km²`
        : 'World land',
  },
)

import { learningPoverty } from '@charts-poc/demo-data/learning-poverty'
import { geoCentroid } from 'd3-geo'
import { worldCountries } from '../fixtures/country-atlas'
import type { LearningPovertyRow } from '@charts-poc/demo-data/learning-poverty'
import type { ExtendedFeature, GeoGeometryObjects } from 'd3-geo'
import type { CountryFeature, CountryGeometry } from '../fixtures/country-atlas'

type PointGeometry = Extract<GeoGeometryObjects, { type: 'Point' }>

export interface LearningPovertyProperties extends LearningPovertyRow {
  name: string
}

export type LearningPovertyCountry = ExtendedFeature<
  CountryGeometry,
  LearningPovertyProperties
>
export type LearningPovertyPoint = ExtendedFeature<
  PointGeometry,
  LearningPovertyProperties
>

// The source uses World Bank names while world-atlas uses Natural Earth names.
// Tiny states absent from the 110m atlas remain unmatched.
const naturalEarthNameBySourceName: Readonly<Record<string, string>> = {
  'Congo, Dem Rep': 'Dem. Rep. Congo',
  'Congo, Rep': 'Congo',
  'Cote d’Ivoire': "Côte d'Ivoire",
  'Czech Republic': 'Czechia',
  'Dominican Republic': 'Dominican Rep.',
  'Egypt, Arab Rep': 'Egypt',
  'Iran, Islamic Rep': 'Iran',
  'Korea, Rep': 'South Korea',
  'Kyrgyz Republic': 'Kyrgyzstan',
  'Russian Federation': 'Russia',
  'Slovak Republic': 'Slovakia',
  'United States': 'United States of America',
  'Yemen, Rep': 'Yemen',
}

const countryByName = new Map(
  worldCountries.map((country) => [country.properties.name, country]),
)

export const learningPovertyCountries: readonly LearningPovertyCountry[] =
  learningPoverty.flatMap((row) => {
    const sourceName = row['Country Name']
    const atlasName = naturalEarthNameBySourceName[sourceName] ?? sourceName
    const country = countryByName.get(atlasName)
    return country ? [joinCountry(country, row)] : []
  })

if (learningPovertyCountries.length !== 95) {
  throw new TypeError(
    `Expected 95 learning-poverty countries in world-atlas, got ${learningPovertyCountries.length}`,
  )
}

export const learningPovertyPoints: readonly LearningPovertyPoint[] =
  learningPovertyCountries.map((country) => ({
    type: 'Feature',
    id: country.id,
    geometry: {
      type: 'Point',
      coordinates: geoCentroid(country),
    },
    properties: country.properties,
  }))

// Largest symbols render first so smaller countries remain selectable.
export const learningPovertyPointsByPopulation: readonly LearningPovertyPoint[] =
  [...learningPovertyPoints].sort(
    (left, right) => right.properties.population - left.properties.population,
  )

function joinCountry(
  country: CountryFeature,
  row: LearningPovertyRow,
): LearningPovertyCountry {
  return {
    type: 'Feature',
    id: country.id,
    geometry: country.geometry,
    properties: {
      name: country.properties.name,
      ...row,
    },
  }
}

import countriesAtlasJson from 'world-atlas/countries-110m.json'
import landAtlasJson from 'world-atlas/land-110m.json'
import detailedLandAtlasJson from 'world-atlas/land-50m.json'
import { geoGraticule10 } from 'd3-geo'
import { feature } from 'topojson-client'
import type {
  ExtendedFeature,
  ExtendedFeatureCollection,
  GeoGeometryObjects,
  GeoSphere,
} from 'd3-geo'

type AtlasTopology = Parameters<typeof feature>[0]

export type CountryGeometry = Extract<
  GeoGeometryObjects,
  { type: 'Polygon' | 'MultiPolygon' }
>

export interface CountryProperties {
  name: string
}

export type CountryFeature = ExtendedFeature<CountryGeometry, CountryProperties>
export type LandFeature = ExtendedFeature<CountryGeometry, Record<never, never>>

export const worldSphere: GeoSphere = { type: 'Sphere' }
export const worldGraticule = geoGraticule10()

const countriesTopology = atlasTopology(
  countriesAtlasJson,
  'world-atlas countries-110m',
)
const countriesObject = countriesTopology.objects.countries
if (!countriesObject) {
  throw new TypeError('world-atlas countries-110m is missing countries')
}

const convertedCountries = feature(countriesTopology, countriesObject)
if (convertedCountries.type !== 'FeatureCollection') {
  throw new TypeError('world-atlas countries did not produce a collection')
}

export const worldCountries: readonly CountryFeature[] =
  convertedCountries.features.flatMap<CountryFeature>((entry) => {
    if (
      !isCountryGeometry(entry.geometry) ||
      !isRecord(entry.properties) ||
      typeof entry.properties.name !== 'string'
    ) {
      return []
    }

    return [
      {
        type: 'Feature',
        id: entry.id === undefined ? entry.properties.name : String(entry.id),
        geometry: entry.geometry,
        properties: {
          name: entry.properties.name,
        },
      },
    ]
  })

if (worldCountries.length !== 177) {
  throw new TypeError(
    `Expected 177 world-atlas countries, got ${worldCountries.length}`,
  )
}

export const worldCountryCollection: ExtendedFeatureCollection<CountryFeature> =
  {
    type: 'FeatureCollection',
    features: [...worldCountries],
  }

export const worldLand = convertLand(landAtlasJson, 'world-atlas land-110m')
export const detailedWorldLand = convertLand(
  detailedLandAtlasJson,
  'world-atlas land-50m',
)

function atlasTopology(value: unknown, label: string): AtlasTopology {
  if (
    !isRecord(value) ||
    value.type !== 'Topology' ||
    !Array.isArray(value.arcs) ||
    !isRecord(value.objects)
  ) {
    throw new TypeError(`${label} is not valid TopoJSON`)
  }
  return value as unknown as AtlasTopology
}

function convertLand(value: unknown, label: string): LandFeature {
  const topology = atlasTopology(value, label)
  const landObject = topology.objects.land
  if (!landObject) {
    throw new TypeError(`${label} is missing land`)
  }

  const converted = feature(topology, landObject)
  const land =
    converted.type === 'FeatureCollection' ? converted.features[0] : converted
  if (!land || land.type !== 'Feature' || !isCountryGeometry(land.geometry)) {
    throw new TypeError(`${label} did not produce polygon geometry`)
  }

  return {
    type: 'Feature',
    geometry: land.geometry,
    properties: {},
  }
}

function isCountryGeometry(
  geometry: GeoGeometryObjects,
): geometry is CountryGeometry {
  return geometry.type === 'Polygon' || geometry.type === 'MultiPolygon'
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === 'object' && value !== null
}