import { useMemo } from "react";
import * as d3 from "d3";

const MARGIN = { top: 50, right: 50, bottom: 50, left: 50 };

type HeatmapProps = {
  width: number;
  height: number;
  data: { x: number; y: string; group: string; value: number }[];
};
  data = [
  { group: "A", y: "Apple", x: 0, value: 68 },
  { group: "A", y: "Orange", x: 0, value: 45 },
  { group: "A", y: "Grapes", x: 0, value: 43 },
  { group: "A", y: "Apple", x: 1, value: 58 },
  { group: "A", y: "Orange", x: 1, value: 85 },
  { group: "A", y: "Grapes", x: 1, value: 23 },
  { group: "A", y: "Apple", x: 2, value: 58 },
  { group: "A", y: "Orange", x: 2, value: 45 },
  { group: "A", y: "Grapes", x: 2, value: 53 },
  { group: "A", y: "Apple", x: 3, value: 58 },
  { group: "A", y: "Orange", x: 3, value: 45 },
  { group: "A", y: "Grapes", x: 3, value: 53 },
  { group: "B", y: "Apple", x: 1, value: 70 },
  { group: "B", y: "Orange", x: 1, value: 52 },
  { group: "B", y: "Grapes", x: 1, value: 65 },
  { group: "B", y: "Apple", x: 2, value: 70 },
  { group: "B", y: "Orange", x: 2, value: 52 },
  { group: "B", y: "Grapes", x: 2, value: 65 },
  { group: "B", y: "Apple", x: 3, value: 70 },
  { group: "B", y: "Orange", x: 3, value: 52 },
  { group: "B", y: "Grapes", x: 3, value: 65 },
  { group: "C", y: "Apple", x: 1, value: 48 },
  { group: "C", y: "Orange", x: 1, value: 46 },
  { group: "C", y: "Grapes", x: 1, value: 65 },
  { group: "C", y: "Apple", x: 2, value: 48 },
  { group: "C", y: "Orange", x: 2, value: 46 },
  { group: "C", y: "Grapes", x: 2, value: 65 },
  { group: "C", y: "Apple", x: 3, value: 48 },
  { group: "C", y: "Orange", x: 3, value: 46 },
  { group: "C", y: "Grapes", x: 3, value: 65 },
  { group: "C", y: "Apple", x: 4, value: 48 },
  { group: "C", y: "Orange", x: 4, value: 46 },
  { group: "C", y: "Grapes", x: 4, value: 65 },
  { group: "C", y: "Apple", x: 5, value: 48 },
  { group: "C", y: "Orange", x: 5, value: 46 },
  { group: "C", y: "Grapes", x: 5, value: 65 },
];

export const Heatmap = ({ width, height, data }: HeatmapProps) => {
  // bounds = area inside the axis
  const boundsWidth = width - MARGIN.right - MARGIN.left;
  const boundsHeight = height - MARGIN.top - MARGIN.bottom;

  // Groups
  const yData = useMemo(() => [...new Set(data.map((d) => d.y))], [data]);
  const xGroups = useMemo(() => [...new Set(data.map((d) => d.group))], [data]);

  const [min, max] = d3.extent(data.map((d) => d.value));

  if (!min || !max) {
    return null;
  }

  // Color scale
  const colorScale = d3
    .scaleSequential()
    .interpolator(d3.interpolateInferno)
    .domain([min, max]);

  // Y scale
  const yScale = d3
    .scaleBand()
    .domain(yData)
    .range([0, boundsHeight])
    .padding(0.01);

  // Group elements
  const groupElements = xGroups.map((group, i) => {
    const groupData = data.filter((d) => d.group === group);
    const days = [...new Set(groupData.map((d) => d.x))];

    // X scale for the group
    const xScale = d3
      .scaleBand()
      .range([0, boundsWidth])
      .domain(days)
      .padding(0.01);

    // Rectangles for the group
    const rects = days.map((day, j) => {
      const dayData = groupData.filter((d) => d.x === day);
      return dayData.map((d, k) => (
        <rect
          key={`rect-${group}-${day}-${k}`}
          x={xScale(day)}
          y={yScale(d.y)}
          width={xScale.bandwidth()}
          height={yScale.bandwidth()}
          fill={colorScale(d.value)}
          rx={5}
          stroke={"white"}
        />
      ));
    });

    // X labels for the group
    const xLabels = days.map((day, j) => (
      <text
        key={`${group}-${day}`}
        x={xScale(day) + xScale.bandwidth() / 2}
        y={boundsHeight + MARGIN.bottom / 2}
        textAnchor="middle"
        dominantBaseline="middle"
        fontSize={10}
      >
        {day}
      </text>
    ));

    // Group label
    const groupLabel = (
      <text
        key={`label-${group}`}
        x={boundsWidth / 2}
        y={boundsHeight + MARGIN.bottom - 10}
        textAnchor="middle"
        dominantBaseline="middle"
        fontSize={10}
      >
        {group}
      </text>
    );

    return (
      <g
        key={`group-${group}`}
        transform={`translate(0, ${i * (boundsHeight + MARGIN.bottom)})`}
      >
        {rects}
        {xLabels}
        {groupLabel}
      </g>
    );
  });

  // Y labels
  const yLabels = yData.map((name, i) => (
    <text
      key={i}
      x={-5}
      y={yScale(name) + yScale.bandwidth() / 2}
      textAnchor="end"
      dominantBaseline="middle"
      fontSize={10}
    >
      {name}
    </text>
  ));

  return (
    <div>
      <svg width={width} height={height}>
        <g transform={`translate(${MARGIN.left}, ${MARGIN.top})`}>
          {yLabels}
          {groupElements}
        </g>
      </svg>
    </div>
  );
};

<script src="https://cdnjs.cloudflare.com/ajax/libs/react/18.2.0/umd/react.production.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/react-dom/18.2.0/umd/react-dom.production.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/d3/7.8.5/d3.min.js"></script>