import { colorLegend, defineChart, rect, text } from '@tanstack/charts'
import { format } from 'd3-format'
import { scaleLinear } from 'd3-scale'
import { survey } from '@charts-poc/demo-data/survey'
import { mosaicLayout, mosaicResponses } from './layout'
import { tanstackMount } from '../../shared/mount'

const percent = format('.0%')
const colors = ['#991b1b', '#ef4444', '#cbd5e1', '#60a5fa', '#1d4ed8']

const definition = () => {
  const layout = mosaicLayout(survey)

  return defineChart({
    marks: [
      rect(layout.cells, {
        x1: 'x1',
        x2: 'x2',
        y1: 'y1',
        y2: 'y2',
        color: 'Response',
        inset: 1,
      }),
      text(layout.labels, {
        x: 'x',
        y: 'y',
        text: 'Question',
        fill: '#334155',
        fontSize: 11,
      }),
    ],
    x: {
      scale: scaleLinear().domain([0, 1]),
      axis: { ticks: { format: percent }, label: 'Share of responses' },
    },
    y: {
      scale: scaleLinear().domain([0, 1.12]),
      axis: { ticks: { format: percent }, label: 'Within-question share' },
    },
    color: {
      domain: mosaicResponses,
      range: colors,
      legend: colorLegend({ label: 'Response' }),
    },
  })
}

export const mount = tanstackMount(definition, 'Marimekko survey composition')

import { cumsum, group, rollup, sum } from 'd3-array'
import type { SurveyRow } from '@charts-poc/demo-data/survey'

export const mosaicResponses = [
  'Strongly Disagree',
  'Disagree',
  'Neutral',
  'Agree',
  'Strongly Agree',
] as const

export interface MosaicCell {
  Question: string
  Response: string
  count: number
  x1: number
  x2: number
  y1: number
  y2: number
}

export interface MosaicLabel {
  Question: string
  x: number
  y: number
}

export function mosaicLayout(rows: readonly SurveyRow[]): {
  cells: readonly MosaicCell[]
  labels: readonly MosaicLabel[]
} {
  const byQuestion = group(rows, (row) => row.Question)
  const questions = Array.from(byQuestion.keys())
  const totals = questions.map(
    (question) => byQuestion.get(question)?.length ?? 0,
  )
  const grandTotal = sum(totals)
  if (grandTotal === 0) return { cells: [], labels: [] }

  const xEnds = cumsum(totals, (count) => count / grandTotal)
  const cells: MosaicCell[] = []
  const labels: MosaicLabel[] = []

  questions.forEach((Question, questionIndex) => {
    const questionRows = byQuestion.get(Question) ?? []
    const questionTotal = totals[questionIndex] ?? 0
    if (questionTotal === 0) return

    const responseCounts = rollup(
      questionRows,
      (values) => values.length,
      (row) => row.Response,
    )
    const x1 = questionIndex === 0 ? 0 : (xEnds[questionIndex - 1] ?? 0)
    const x2 = xEnds[questionIndex] ?? x1
    const ordered = mosaicResponses.map((Response) => ({
      Question,
      Response,
      count: responseCounts.get(Response) ?? 0,
    }))
    const yEnds = cumsum(ordered, (row) => row.count / questionTotal)

    ordered.forEach((row, responseIndex) => {
      cells.push({
        ...row,
        x1,
        x2,
        y1: responseIndex === 0 ? 0 : (yEnds[responseIndex - 1] ?? 0),
        y2: yEnds[responseIndex] ?? 0,
      })
    })
    labels.push({
      Question,
      x: (x1 + x2) / 2,
      y: 1.055,
    })
  })

  return { cells, labels }
}