import { barX, colorLegend, defineChart, ruleX } from '@tanstack/charts'
import { survey } from '@charts-poc/demo-data/survey'
import { scaleBand, scaleLinear } from 'd3-scale'
import { likertResponses, selectLikertSurvey } from './selection'
import { likertSegments } from './transform'
import { tanstackMount } from '../../shared/mount'

const colors = ['#991b1b', '#ef4444', '#cbd5e1', '#60a5fa', '#1d4ed8']

const likertSurvey = selectLikertSurvey(survey)
const segments = likertSegments(likertSurvey)

const definition = () =>
  defineChart({
    marks: [
      barX(segments, {
        x1: 'x1',
        x2: 'x2',
        y: 'Question',
        color: 'Response',
        key: (row) => `${row.Question}:${row.Response}`,
        inset: 0.75,
      }),
      ruleX([0], { stroke: '#64748b' }),
    ],
    x: {
      scale: scaleLinear,
      grid: true,
      axis: {
        ticks: { format: (value) => `${Math.abs(value)}` },
        label: '← more disagree · Number of responses · more agree →',
      },
    },
    y: {
      scale: () => scaleBand<string>().paddingInner(0.14).paddingOuter(0.08),
    },
    color: {
      domain: likertResponses,
      range: colors,
      legend: colorLegend({ label: 'Response' }),
    },
  })

export const mount = tanstackMount(
  definition,
  'Diverging Likert survey responses',
)

import type { SurveyRow } from '@charts-poc/demo-data/survey'

export type LikertResponse =
  'Strongly Disagree' | 'Disagree' | 'Neutral' | 'Agree' | 'Strongly Agree'

export interface LikertSurveyRow extends SurveyRow {
  readonly Response: LikertResponse
}

export const likertResponses: readonly LikertResponse[] = [
  'Strongly Disagree',
  'Disagree',
  'Neutral',
  'Agree',
  'Strongly Agree',
]

const responseSet = new Set<string>(likertResponses)

export function selectLikertSurvey(
  rows: readonly SurveyRow[],
): readonly LikertSurveyRow[] {
  return rows.filter((row): row is LikertSurveyRow =>
    responseSet.has(row.Response),
  )
}

export function likertQuestions(rows: readonly LikertSurveyRow[]) {
  return [...new Set(rows.map((row) => row.Question))]
}

import { group } from 'd3-array'
import { likertResponses } from './selection'
import type { LikertResponse, LikertSurveyRow } from './selection'

export interface LikertSegment {
  readonly Question: string
  readonly Response: LikertResponse
  readonly count: number
  readonly x1: number
  readonly x2: number
}

export function likertSegments(
  rows: readonly LikertSurveyRow[],
): readonly LikertSegment[] {
  return [...group(rows, (row) => row.Question)].flatMap(
    ([Question, questionRows]) => {
      const counts = new Map<LikertResponse, number>(
        likertResponses.map((Response) => [
          Response,
          questionRows.filter((row) => row.Response === Response).length,
        ]),
      )
      let cursor =
        -(counts.get('Strongly Disagree') ?? 0) -
        (counts.get('Disagree') ?? 0) -
        (counts.get('Neutral') ?? 0) / 2

      return likertResponses.map((Response) => {
        const count = counts.get(Response) ?? 0
        const x1 = cursor
        cursor += count
        return { Question, Response, count, x1, x2: cursor }
      })
    },
  )
}