Scatter Marginal Histograms

import { type ChartPoint, binX, binY, defineChart, dot, rect } from "@tanstack/charts";
import { motion } from "@tanstack/charts/motion";
import { tooltip } from "@tanstack/charts/tooltip";
import { viewGrid } from "@tanstack/charts/view";
import { Chart } from "@tanstack/react-charts/core";
import { useMemo } from "react";

import { scaleLinear } from "d3-scale";
import { BarChart3 } from "lucide-react";

import { Card, CardContent, CardDescription, CardFooter, CardHeader, CardTitle } from "@/components/ui/card";

interface PenguinDatum {
    species: "Adelie" | "Chinstrap" | "Gentoo";
    flipper_length_mm: number;
    body_mass_g: number;
}

const motionRenderer = motion({
    initial: "always",
    transition: { type: "spring", stiffness: 170, damping: 18, mass: 1 },
});

const flipperBoundaries = [170, 180, 190, 200, 210, 220, 230, 240];
const massBoundaries = [2500, 3000, 3500, 4000, 4500, 5000, 5500, 6000, 6500];

function generateDemoData(): PenguinDatum[] {
    const results: PenguinDatum[] = [];
    const speciesList: ("Adelie" | "Chinstrap" | "Gentoo")[] = ["Adelie", "Chinstrap", "Gentoo"];

    speciesList.forEach((sp) => {
        let fMean = 190;
        let mMean = 3700;
        if (sp === "Chinstrap") {
            fMean = 195;
            mMean = 3730;
        }
        if (sp === "Gentoo") {
            fMean = 217;
            mMean = 5080;
        }

        for (let i = 0; i < 35; i++) {
            const flipper = fMean + (Math.random() - 0.5) * 16;
            const mass = mMean + (Math.random() - 0.5) * 900;
            results.push({
                species: sp,
                flipper_length_mm: Math.round(flipper),
                body_mass_g: Math.round(mass),
            });
        }
    });

    return results;
}

function shadcnTooltipContent<TDatum>(points: readonly ChartPoint<TDatum>[]) {
    if (points.length === 0) return { title: "", rows: [] };
    const d = points[0].datum as PenguinDatum | { x1?: number; x2?: number; y1?: number; y2?: number; count?: number };
    if (!d) return { title: "", rows: [] };

    if ("count" in d && "x1" in d && d.x1 !== undefined) {
        return {
            title: `Flipper Range: ${d.x1}–${d.x2} mm`,
            rows: [{ label: "Penguin Count", value: `${d.count}`, color: "var(--chart-1)" }],
        };
    }
    if ("count" in d && "y1" in d && d.y1 !== undefined) {
        return {
            title: `Mass Range: ${d.y1}–${d.y2} g`,
            rows: [{ label: "Penguin Count", value: `${d.count}`, color: "var(--chart-2)" }],
        };
    }
    const p = d as PenguinDatum;
    return {
        title: `${p.species} Penguin`,
        rows: [
            { label: "Flipper Length", value: `${p.flipper_length_mm} mm`, color: "var(--chart-1)" },
            { label: "Body Mass", value: `${p.body_mass_g} g`, color: "var(--chart-2)" },
        ],
    };
}

export function ScatterMarginalHistograms() {
    const data = useMemo(() => generateDemoData(), []);

    const xBins = useMemo(
        () =>
            binX(data, {
                value: "flipper_length_mm",
                thresholds: flipperBoundaries,
                outputs: { count: { reduce: "count" } },
            }),
        [data],
    );

    const yBins = useMemo(
        () =>
            binY(data, {
                value: "body_mass_g",
                thresholds: massBoundaries,
                outputs: { count: { reduce: "count" } },
            }),
        [data],
    );

    const flipperScale = useMemo(() => scaleLinear().domain([170, 240]), []);
    const massScale = useMemo(() => scaleLinear().domain([2500, 6500]), []);

    const definition = useMemo(
        () =>
            defineChart(
                viewGrid({
                    id: "penguin-marginals",
                    rows: [
                        { id: "top", size: 60 },
                        { id: "main", grow: 1 },
                    ],
                    columns: [
                        { id: "main", grow: 1 },
                        { id: "right", size: 60 },
                    ],
                    gap: 6,
                    views: [
                        {
                            id: "main",
                            row: "main",
                            column: "main",
                            chart: defineChart({
                                marks: [
                                    dot(data, {
                                        id: "penguins",
                                        x: "flipper_length_mm",
                                        y: "body_mass_g",
                                        color: "species",
                                        r: 3.5,
                                        fillOpacity: 0.85,
                                    }),
                                ],
                                scales: {
                                    x: {
                                        scale: flipperScale,
                                        grid: true,
                                        axis: { label: "Flipper Length (mm)" },
                                    },
                                    y: {
                                        scale: massScale,
                                        grid: true,
                                        axis: { label: "Body Mass (g)" },
                                    },
                                },
                                color: {
                                    range: ["var(--chart-1)", "var(--chart-2)", "var(--chart-3)"],
                                },
                            }),
                        },
                        {
                            id: "top",
                            row: "top",
                            column: "main",
                            share: { x: "main" },
                            chart: defineChart({
                                marks: [
                                    rect(xBins, {
                                        id: "flipper-histogram",
                                        x: "x",
                                        x1: "x1",
                                        x2: "x2",
                                        y: "count",
                                        y1: () => 0,
                                        y2: "count",
                                        fill: "var(--chart-1)",
                                        fillOpacity: 0.75,
                                        inset: 1,
                                    }),
                                ],
                                scales: {
                                    x: { scale: flipperScale },
                                    y: { scale: scaleLinear },
                                },
                                guides: false,
                            }),
                        },
                        {
                            id: "right",
                            row: "main",
                            column: "right",
                            share: { y: "main" },
                            chart: defineChart({
                                marks: [
                                    rect(yBins, {
                                        id: "mass-histogram",
                                        x: "count",
                                        x1: () => 0,
                                        x2: "count",
                                        y: "y",
                                        y1: "y1",
                                        y2: "y2",
                                        fill: "var(--chart-2)",
                                        fillOpacity: 0.75,
                                        inset: 1,
                                    }),
                                ],
                                scales: {
                                    x: { scale: scaleLinear },
                                    y: { scale: massScale },
                                },
                                guides: false,
                            }),
                        },
                    ],
                }),
                {
                    keyboard: true,
                    tooltip: {
                        use: tooltip,
                        className:
                            "font-sans shadow-none rounded-md bg-popover text-popover-foreground border p-2 text-xs",
                        anchor: "pointer",
                        placement: "auto",
                        content: (points) => shadcnTooltipContent(points),
                    },
                },
            ),
        [data, xBins, yBins, flipperScale, massScale],
    );

    return (
        <Card className="w-full max-w-xl">
            <CardHeader>
                <CardTitle>Scatterplot with Marginal Histograms</CardTitle>
                <CardDescription>
                    Bivariate scatter paired with top and right 1D marginal frequency bins
                </CardDescription>
            </CardHeader>
            <CardContent>
                <div className="h-[280px] w-full">
                    <Chart
                        definition={definition}
                        renderer={motionRenderer}
                        height={280}
                        ariaLabel="Scatterplot with Marginal Histograms Chart"
                    />
                </div>
            </CardContent>
            <CardFooter className="flex-col items-start gap-2 text-sm">
                <div className="flex items-center gap-2 leading-none font-medium">
                    Marginal bars summarize 1D distribution peaks alongside 2D point cloud{" "}
                    <BarChart3 className="h-4 w-4 text-emerald-500" />
                </div>
                <div className="text-muted-foreground leading-none">
                    Multi-view grid layout aligns top X-histogram and right Y-histogram seamlessly
                </div>
            </CardFooter>
        </Card>
    );
}

export default ScatterMarginalHistograms;