import { simpsons } from '@charts-poc/demo-data/simpsons'
import { areaY, d3Curve, defineChart, lineY, ruleY } from '@tanstack/charts'
import { scaleLinear } from 'd3-scale'
import { curveBasis } from 'd3-shape'
import { isRatedEpisode, ridgeDensity, ridgeSeasons } from './transform'
import { tanstackMount } from '../../shared/mount'
import type { ConformanceInput } from '../../types'

const colors = ['#2563eb', '#0d9488', '#d97706']

const definition = (input: ConformanceInput) => {
  const seasons = ridgeSeasons(input.revision)
  const rows = ridgeDensity(simpsons.filter(isRatedEpisode), seasons)
  const curve = d3Curve(curveBasis)

  return defineChart({
    marks: [
      ruleY([0, 1, 2], {
        stroke: '#94a3b8',
        strokeOpacity: 0.5,
      }),
      areaY(rows, {
        x: 'imdb_rating',
        y1: 'baseline',
        y2: 'density',
        color: 'season',
        fillOpacity: 0.52,
        curve,
      }),
      lineY(rows, {
        x: 'imdb_rating',
        y: 'density',
        color: 'season',
        strokeWidth: 1.5,
        curve,
      }),
    ],
    x: {
      scale: scaleLinear().domain([4, 10]),
      grid: true,
      axis: { label: 'IMDb rating' },
    },
    y: {
      scale: scaleLinear().domain([-0.08, 2.86]),
      axis: {
        ticks: {
          count: seasons.length,
          format: (value) => {
            const season = seasons[Math.round(value)]
            return season === undefined ? '' : `Season ${season}`
          },
        },
      },
    },
    color: {
      range: colors,
    },
    margin: { left: 76 },
  })
}

export const mount = tanstackMount(definition, 'Ridgeline density comparison')

import { bin } from 'd3-array'
import type { SimpsonsRow } from '@charts-poc/demo-data/simpsons'

export type RatedEpisode = SimpsonsRow & {
  readonly imdb_rating: number
}

export interface RidgePoint {
  id: string
  season: number
  imdb_rating: number
  baseline: number
  density: number
}

const boundaries = [
  4, 4.25, 4.5, 4.75, 5, 5.25, 5.5, 5.75, 6, 6.25, 6.5, 6.75, 7, 7.25, 7.5,
  7.75, 8, 8.25, 8.5, 8.75, 9, 9.25, 9.5, 9.75, 10,
]
const createBins = bin<RatedEpisode, number>()
  .value((row) => row.imdb_rating)
  .domain([4, 10])
  .thresholds(boundaries.slice(1, -1))

export function isRatedEpisode(row: SimpsonsRow): row is RatedEpisode {
  return row.imdb_rating !== null
}

export function ridgeSeasons(revision: number): readonly number[] {
  const offset = revision % 2
  return [1 + offset, 10 + offset, 20 + offset]
}

export function ridgeDensity(
  episodes: readonly RatedEpisode[],
  seasons: readonly number[],
): readonly RidgePoint[] {
  return seasons.flatMap((season, seasonIndex) => {
    const buckets = createBins(
      episodes.filter((episode) => episode.season === season),
    )
    const maximum = Math.max(...buckets.map((bucket) => bucket.length), 1)

    return buckets.flatMap((bucket, index) => {
      if (bucket.x0 === undefined || bucket.x1 === undefined) return []
      return [
        {
          id: `${season}:${index}`,
          season,
          imdb_rating: (bucket.x0 + bucket.x1) / 2,
          baseline: seasonIndex,
          density: seasonIndex + (bucket.length / maximum) * 0.78,
        },
      ]
    })
  })
}