// SVG setup
const svg = d3.select("svg"),
      margin = {top: 80, right: 30, bottom: 150, left: 160},
      width = +svg.attr("width") - margin.left - margin.right,
      height = +svg.attr("height") - margin.top - margin.bottom,
      g = svg.append("g").attr("transform", `translate(${margin.left},${margin.top})`);

// Scales for grouped bar and dot plots
const x0 = d3.scaleBand().rangeRound([0, width]).paddingInner(0.1);
const x1 = d3.scaleBand().padding(0.05);
const y = d3.scaleLinear().rangeRound([height, 0]);

// Education levels for legend and grouped charts
const educationLevels = [
  "Primary Education",
  "Secondary Education",
  "Tertiary Education"
];

// Color scale for education levels
const color = d3.scaleOrdinal()
  .domain(educationLevels)
  .range(["#4E79A7", "#F28E2B", "#59A14F"]);

// Tooltip div
const tooltip = d3.select("body").append("div").attr("class", "tooltip");

// Chart titles
const chartTitle1 = svg.append("text")
  .attr("x", +svg.attr("width") / 2)
  .attr("y", 30)
  .attr("text-anchor", "middle")
  .style("font-size", "20px")
  .style("font-weight", "bold");

const chartTitle2 = svg.append("text")
  .attr("x", +svg.attr("width") / 2)
  .attr("y", 55)
  .attr("text-anchor", "middle")
  .style("font-size", "18px")
  .style("fill", "#666");

// "No Data" placeholder text
const noDataText = g.append("text")
  .attr("x", width / 2)
  .attr("y", height / 2)
  .attr("text-anchor", "middle")
  .style("font-size", "18px")
  .style("fill", "#999")
  .style("display", "none")
  .text("No data available for selected filters");

// Dark mode toggle button
const darkModeToggle = document.createElement('button');
darkModeToggle.textContent = 'Toggle Dark Mode';
Object.assign(darkModeToggle.style, {
  margin: '5px',
  padding: '8px 12px',
  fontSize: '14px',
  cursor: 'pointer',
  border: 'none',
  borderRadius: '4px',
  backgroundColor: '#333',
  color: '#fff'
});
document.querySelector('.controls').appendChild(darkModeToggle);

// Dark mode logic
let darkMode = false;
darkModeToggle.addEventListener('click', () => {
  darkMode = !darkMode;
  document.body.style.backgroundColor = darkMode ? '#121212' : '#fff';
  document.body.style.color = darkMode ? '#eee' : '#000';
  chartTitle1.style("fill", darkMode ? "#eee" : "#000");
  chartTitle2.style("fill", darkMode ? "#ccc" : "#666");
  darkModeToggle.style.backgroundColor = darkMode ? '#eee' : '#333';
  darkModeToggle.style.color = darkMode ? '#000' : '#fff';
  tooltip.style("background-color", darkMode ? "#222" : "#fff")
         .style("color", darkMode ? "#eee" : "#000");
  g.selectAll("text").style("fill", darkMode ? "#eee" : "#000");
  g.selectAll(".lowest-annotation").style("fill", darkMode ? "#eee" : "red");
  g.selectAll(".lowest-box").attr("fill", darkMode ? "#222" : "#fff").attr("stroke", darkMode ? "#aaa" : "#999");
  // Update the chart as different colours are used in different charts (namely the choropleth)
  updateChart();
});


// Add 'choropleth' to chart type selector on page load
if (document.getElementById("chartTypeSelect") && !Array.from(document.getElementById("chartTypeSelect").options).some(opt => opt.value === "choropleth")) {
  const option = document.createElement("option");
  option.value = "choropleth";
  option.text = "Choropleth Map";
  document.getElementById("chartTypeSelect").appendChild(option);
}

// Load both CSV and TopoJSON data
Promise.all([
  d3.csv("PHSwithContinent.csv"),
  d3.json("script/countries-50m.json")
]).then(function([data, world]) {
  // Normalize Socio-economic status values to match educationLevels
  data.forEach(d => {
    if (d["Socio-economic status"] === "Pre-primary, primary and lower secondary education") {
      d["Socio-economic status"] = "Primary Education";
    } else if (d["Socio-economic status"] === "Upper secondary and post-secondary non-tertiary, all programmes") {
      d["Socio-economic status"] = "Secondary Education";
    } else if (d["Socio-economic status"] === "Tertiary education") {
      d["Socio-economic status"] = "Tertiary Education";
    }
  });
  const allData = data;
  const worldData = world;

  // Populate continent and education dropdowns
  const continents = Array.from(new Set(data.map(d => d.CONTINENT))).sort();
  continents.forEach(continent => {
    d3.select("#continentSelect")
      .append("option")
      .attr("value", continent)
      .text(continent);
  });

  educationLevels.forEach(level => {
    d3.select("#incomeSelect")
      .append("option")
      .attr("value", level)
      .text(level);
  });

  // Setup static legend (initial)
  const legend = d3.select("#legend");
  color.domain().forEach(d => {
    const item = legend.append("div").attr("class", "legend-item");
    item.append("div")
      .attr("class", "legend-color")
      .style("background-color", color(d));
    item.append("span").text(d);
  });

  // Update the chart
  function updateChart() {
    const selectedContinent = d3.select("#continentSelect").property("value");
    const selectedSex = d3.select("#sexSelect").property("value");
    const selectedEdu = d3.select("#incomeSelect").property("value");
    const selectedChart = d3.select("#chartTypeSelect").property("value");

    updateLegend(selectedChart);

    // Remove zoom for non-choropleth charts
    if (selectedChart !== "choropleth") {
      svg.on(".zoom", null); // Remove zoom event listeners
      g.attr("transform", `translate(${margin.left},${margin.top})`); // Reset transform
    }

    // Update chart titles
    let title1 = "Self-Reported Health";
    let title2 = "";
    if (selectedContinent !== "All") title2 += `${selectedContinent}`;
    if (selectedSex !== "All") title2 += (title2 ? ", " : "") + selectedSex;
    if (selectedEdu !== "All") title2 += (title2 ? ", " : "") + selectedEdu;
    chartTitle1.text(title1);
    chartTitle2.text(title2);

    // Filter data based on selections
    let filteredData = allData.filter(d =>
      (selectedContinent === "All" || d.CONTINENT === selectedContinent) &&
      (selectedSex === "All" || d.Sex === selectedSex) &&
      (selectedEdu === "All" || d["Socio-economic status"] === selectedEdu)
    ).filter(d =>
      educationLevels.includes(d["Socio-economic status"]) &&
      !isNaN(parseFloat(d.OBS_VALUE))
    );

    // Clear and redraw
    g.selectAll("*").remove();
    if (filteredData.length === 0) {
      noDataText.style("display", null);
      return;
    } else {
      noDataText.style("display", "none");
    }

    if (selectedChart === "grouped") {
      drawGrouped(filteredData);
    } else if (selectedChart === "dot") {
      drawDotPlot(filteredData);
    } else if (selectedChart === "heatmap") {
      drawHeatmap(filteredData);
    } else if (selectedChart === "choropleth") {
      drawChoropleth(filteredData, worldData, selectedContinent);
    }
  }

  updateChart();
  d3.select("#loading").classed("hidden", true);

  // Attach event listeners
  d3.selectAll("select").on("change", updateChart);
  d3.select("#resetButton").on("click", function() {
    d3.select("#continentSelect").property("value", "All");
    d3.select("#sexSelect").property("value", "All");
    d3.select("#incomeSelect").property("value", "All");
    d3.select("#chartTypeSelect").property("value", "grouped");
    updateChart();
  });

  d3.select("#downloadButton").on("click", function() {
    const c = d3.select("#continentSelect").property("value");
    const s = d3.select("#sexSelect").property("value");
    const e = d3.select("#incomeSelect").property("value");
    const chartType = d3.select("#chartTypeSelect").property("value");
  
    const parts = [
      "education_health",
      chartType,                      
      c !== "All" ? c : null,
      s !== "All" ? s : null,
      e !== "All" ? e.replace(/[^a-zA-Z0-9]+/g, "_").toLowerCase() : null
    ].filter(Boolean);
  
    const filename = parts.join("_") + ".png";
  
    saveSvgAsPng(document.getElementById("chart"), filename, {
      scale: 2,
      backgroundColor: darkMode ? "#121212" : "#ffffff"
    });
  });

  // Add Export to PDF button (reuse downloadButton's class and styles)
  const exportPdfButton = document.createElement('button');
  exportPdfButton.textContent = 'Export to PDF';
  // Copy class from downloadButton if available
  const downloadBtn = document.getElementById("downloadButton");
  if (downloadBtn && downloadBtn.className) {
    exportPdfButton.className = downloadBtn.className;
  }
  // Insert immediately after downloadButton
  if (downloadBtn) {
    downloadBtn.insertAdjacentElement("afterend", exportPdfButton);
  } else {
    // fallback
    document.querySelector('.controls').appendChild(exportPdfButton);
  }

  // PDF export logic
  exportPdfButton.addEventListener('click', async () => {
    // Ensure jsPDF and svg2pdf.js are loaded
    if (!(window.jspdf && window.svg2pdf)) {
      alert('PDF export libraries not loaded.');
      return;
    }
    const { jsPDF } = window.jspdf;
    const svgElement = document.querySelector("svg");
    const svgClone = svgElement.cloneNode(true);
    svgClone.removeAttribute("style");
    // Ensure correct namespace for svg2pdf
    svgClone.setAttribute("xmlns", "http://www.w3.org/2000/svg");

    // Set cloned SVG width and height explicitly
    svgClone.setAttribute("width", svg.attr("width"));
    svgClone.setAttribute("height", svg.attr("height"));

    // Use fixed format and orientation
    const doc = new jsPDF({
      orientation: 'landscape',
      unit: 'pt',
      format: [+svg.attr("width"), +svg.attr("height")]
    });

    // Use standard margin offset and scale
    await window.svg2pdf(svgClone, doc, {
      xOffset: 0,
      yOffset: 0,
      scale: 1
    });

    const c = d3.select("#continentSelect").property("value");
    const s = d3.select("#sexSelect").property("value");
    const e = d3.select("#incomeSelect").property("value");
    const chartType = d3.select("#chartTypeSelect").property("value");

    const parts = [
      "education_health",
      chartType,
      c !== "All" ? c : null,
      s !== "All" ? s : null,
      e !== "All" ? e.replace(/[^a-zA-Z0-9]+/g, "_").toLowerCase() : null
    ].filter(Boolean);
    const filename = parts.join("_") + ".pdf";

    doc.save(filename);
  });
});
// Update legend depending on chart type
function updateLegend(chartType) {
  const legend = d3.select("#legend");
  legend.html("");

  if (chartType === "heatmap") {
    const wrapper = legend.append("div")
      .style("display", "flex")
      .style("flex-direction", "column")
      .style("align-items", "center")
      .style("margin-top", "10px");

    wrapper.append("div")
      .style("width", "300px")
      .style("height", "20px")
      .style("background", "linear-gradient(to right, #edf8b1, #2c7fb8)")
      .style("margin-bottom", "5px");
    wrapper.append("div")
      .style("width", "300px")
      .style("display", "flex")
      .style("justify-content", "space-between")
      .style("font-size", "12px")
      .html('<span>0%</span><span>50%</span><span>100%</span>');
  } else if (chartType === "choropleth") {
    // No drawing of any legend for choropleth
    return;
  } else {
    educationLevels.forEach(level => {
      const item = legend.append("div").attr("class", "legend-item");
      item.append("div")
        .attr("class", "legend-color")
        .style("background-color", color(level));
      item.append("span").text(level);
    });
  }
}

function drawGrouped(data) {
  // Setup x0, x1, and y domains
  const countries = Array.from(new Set(data.map(d => d["Reference area"])));
  x0.domain(countries);
  x1.domain(educationLevels).rangeRound([0, x0.bandwidth()]);
  y.domain([0, 100]);

  // Group data by country
  const groupedData = data.reduce((acc, d) => {
    const found = acc.find(v => v.Country === d["Reference area"]);
    const value = parseFloat(d.OBS_VALUE);
    const entry = { ...d, value };
    if (found) found.values.push(entry);
    else acc.push({ Country: d["Reference area"], values: [entry] });
    return acc;
  }, []);

  // Draw grouped bars
  const countryGroups = g.append("g")
    .selectAll("g")
    .data(groupedData)
    .join("g")
    .attr("transform", d => `translate(${x0(d.Country)},0)`)
    .attr("class", "country-group");

  countryGroups.selectAll("rect")
    .data(d => d.values)
    .join("rect")
    .attr("x", d => x1(d["Socio-economic status"]))
    .attr("width", x1.bandwidth())
    .attr("fill", d => color(d["Socio-economic status"]))
    .attr("y", height)
    .attr("height", 0)
    .on("mouseover", function(event, d) {
      // Show tooltip
      tooltip.html(`<strong>${d["Reference area"]}</strong><br/>${d["Socio-economic status"]}<br/><strong>${d.OBS_VALUE}%</strong>`)
        .style("left", (event.pageX + 10) + "px")
        .style("top", (event.pageY - 40) + "px")
        .transition().duration(300)
        .style("opacity", 1)
        .style("transform", "translateY(-10px)");
    })
    .on("mouseout", function() {
      // Hide tooltip
      tooltip.transition().duration(500)
        .style("opacity", 0)
        .style("transform", "translateY(0px)");
    })
    .transition()
    .duration(800)
    .ease(d3.easeCubicOut)
    .attr("y", d => y(d.value))
    .attr("height", d => height - y(d.value));

  // Add X axis
  g.append("g")
    .attr("transform", `translate(0,${height})`)
    .call(d3.axisBottom(x0))
    .selectAll("text")
    .attr("transform", "rotate(45)")
    .style("text-anchor", "start");

  // Add Y axis
  g.append("g")
    .call(d3.axisLeft(y).ticks(10));

  // Annotate lowest bar
  if (groupedData.length > 0) {
    const flat = groupedData.flatMap(g => g.values);
    const lowest = flat.reduce((min, d) => d.value < min.value ? d : min, flat[0]);

    const barX = x0(lowest["Reference area"]) + x1(lowest["Socio-economic status"]) + x1.bandwidth() / 2;
    const barY = y(lowest.value);

    const labelText = `Lowest: ${lowest["Reference area"]}`;

    // Measure text width
    const tempText = g.append("text")
      .attr("x", -9999)
      .attr("y", -9999)
      .style("font-size", "12px")
      .style("font-weight", "bold")
      .text(labelText);

    const textWidth = tempText.node().getBBox().width;
    tempText.remove();

    // Draw annotation background box
    g.append("rect")
      .attr("class", "lowest-box")
      .attr("x", barX - textWidth / 2 - 6)
      .attr("y", barY - 28)
      .attr("width", textWidth + 12)
      .attr("height", 20)
      .attr("rx", 4)
      .attr("fill", darkMode ? "#222" : "#fff")
      .attr("stroke", darkMode ? "#aaa" : "#999")
      .attr("stroke-width", 0.5)
      .attr("opacity", 0.9);

    // Draw annotation text
    g.append("text")
      .attr("class", "lowest-annotation")
      .attr("x", barX)
      .attr("y", barY - 14)
      .attr("text-anchor", "middle")
      .style("font-size", "12px")
      .style("font-weight", "bold")
      .style("fill", darkMode ? "#eee" : "red")
      .text(labelText);
  }
}

function drawDotPlot(data) {
  // Setup x and y scales
  const x = d3.scaleLinear()
    .domain([0, 100])
    .range([0, width]);

  const y = d3.scaleBand()
    .domain([...new Set(data.map(d => d["Reference area"]))])
    .range([0, height])
    .padding(0.3);

  // Group data by country
  const groupedData = data.reduce((acc, d) => {
    const found = acc.find(v => v.Country === d["Reference area"]);
    const value = parseFloat(d.OBS_VALUE);
    const entry = { ...d, value };
    if (found) found.values.push(entry);
    else acc.push({ Country: d["Reference area"], values: [entry] });
    return acc;
  }, []);

  // Draw dot points
  const countryGroups = g.append("g")
    .selectAll("g")
    .data(groupedData)
    .join("g")
    .attr("transform", d => `translate(0,${y(d.Country)})`)
    .attr("class", "country-group");

  countryGroups.selectAll("circle")
    .data(d => d.values)
    .join("circle")
    .attr("cy", d => 0)
    .attr("cx", 0)
    .attr("r", 0)
    .attr("fill", d => color(d["Socio-economic status"]))
    .on("mouseover", function(event, d) {
      // Show tooltip
      tooltip.html(`<strong>${d["Reference area"]}</strong><br/>${d["Socio-economic status"]}<br/><strong>${d.OBS_VALUE}%</strong>`)
        .style("left", (event.pageX + 10) + "px")
        .style("top", (event.pageY - 40) + "px")
        .transition().duration(300)
        .style("opacity", 1)
        .style("transform", "translateY(-10px)");
    })
    .on("mouseout", function() {
      // Hide tooltip
      tooltip.transition().duration(500)
        .style("opacity", 0)
        .style("transform", "translateY(0px)");
    })
    .transition()
    .duration(800)
    .ease(d3.easeCubicOut)
    .attr("cx", d => x(d.value))
    .attr("r", 6);

  // Add axes
  g.append("g")
    .attr("transform", `translate(0,${height})`)
    .call(d3.axisBottom(x).ticks(10));

  g.append("g")
    .call(d3.axisLeft(y));

  // Add x-axis label
  g.append("text")
    .attr("x", width / 2)
    .attr("y", height + 50)
    .attr("text-anchor", "middle")
    .style("font-size", "14px")
    .style("fill", darkMode ? "#fff" : "#000")
    .text("Percentage Reporting Good Health (%)");

  // Annotate lowest dot
  if (groupedData.length > 0) {
    const flat = groupedData.flatMap(g => g.values);
    const lowest = flat.reduce((min, d) => d.value < min.value ? d : min, flat[0]);

    const dotX = x(lowest.value);
    const dotY = y(lowest["Reference area"]);

    const labelText = `Lowest: ${lowest["Reference area"]}`;

    // Measure text width
    const tempText = g.append("text")
      .attr("x", -9999)
      .attr("y", -9999)
      .style("font-size", "12px")
      .style("font-weight", "bold")
      .text(labelText);

    const textWidth = tempText.node().getBBox().width;
    tempText.remove();

    // Draw annotation background box
    g.append("rect")
      .attr("class", "lowest-box")
      .attr("x", dotX - textWidth / 2 - 6)
      .attr("y", dotY - 28)
      .attr("width", textWidth + 12)
      .attr("height", 20)
      .attr("rx", 4)
      .attr("fill", darkMode ? "#222" : "#fff")
      .attr("stroke", darkMode ? "#aaa" : "#999")
      .attr("stroke-width", 0.5)
      .attr("opacity", 0.9);

    // Draw annotation text
    g.append("text")
      .attr("class", "lowest-annotation")
      .attr("x", dotX)
      .attr("y", dotY - 14)
      .attr("text-anchor", "middle")
      .style("font-size", "12px")
      .style("font-weight", "bold")
      .style("fill", darkMode ? "#eee" : "red")
      .text(labelText);
  }
}

function drawHeatmap(data) {
  // Setup x and y scales for heatmap
  const countries = [...new Set(data.map(d => d["Reference area"]))];

  const x = d3.scaleBand()
    .domain(educationLevels)
    .range([0, width])
    .padding(0.05);

  const y = d3.scaleBand()
    .domain(countries)
    .range([0, height])
    .padding(0.05);

  // Color scale for heatmap (sequential)
  const colorScale = d3.scaleSequential()
    .interpolator(d3.interpolateYlGnBu)
    .domain([0, 100]);

  // Draw heatmap squares
  g.selectAll()
    .data(data)
    .join("rect")
    .attr("x", d => x(d["Socio-economic status"]))
    .attr("y", d => y(d["Reference area"]))
    .attr("width", x.bandwidth())
    .attr("height", y.bandwidth())
    .attr("fill", d => colorScale(parseFloat(d.OBS_VALUE)))
    .on("mouseover", function(event, d) {
      // Show tooltip
      tooltip.html(`<strong>${d["Reference area"]}</strong><br/>${d["Socio-economic status"]}: ${d.OBS_VALUE}%`)
        .style("left", (event.pageX + 10) + "px")
        .style("top", (event.pageY - 40) + "px")
        .transition().duration(300)
        .style("opacity", 1)
        .style("transform", "translateY(-10px)");
    })
    .on("mouseout", function() {
      // Hide tooltip
      tooltip.transition().duration(500)
        .style("opacity", 0)
        .style("transform", "translateY(0px)");
    });

  // Add axes
  g.append("g")
    .attr("transform", `translate(0,${height})`)
    .call(d3.axisBottom(x));

  g.append("g")
    .call(d3.axisLeft(y));
}

// Choropleth map drawing function
function drawChoropleth(data, worldData, selectedContinent) {
  // Remove any previous zoom behavior
  svg.on("mousedown.zoom", null).on("touchstart.zoom", null).on("touchmove.zoom", null).on("touchend.zoom", null);

  // Setup projection and path
  const projection = d3.geoMercator()
    .fitSize([width, height], topojson.feature(worldData, worldData.objects.countries));
  const path = d3.geoPath().projection(projection);

  // Prepare a map from country name to value
  const valueByCountry = {};
  data.forEach(d => {
    valueByCountry[d["Reference area"]] = parseFloat(d.OBS_VALUE);
  });

  // Color scale for choropleth
  const colorScale = d3.scaleSequential()
    .interpolator(darkMode ? d3.interpolateYlOrRd : d3.interpolateYlGnBu)
    .domain(darkMode ? [100, 0] : [0, 100]);

  // Remove previous map if any
  g.selectAll(".country").remove();
  g.selectAll(".choropleth-legend").remove();

  // Draw countries
  const countries = topojson.feature(worldData, worldData.objects.countries).features;
  const countryPaths = g.selectAll(".country")
    .data(countries)
    .join("path")
    .attr("class", "country")
    .attr("d", path)
    .attr("fill", d => {
      const name = d.properties.name;
      const val = valueByCountry[name];
      return val != null ? colorScale(val) : "#ccc";
    })
    .attr("stroke", () => darkMode ? "#fff" : "#000")
    .attr("stroke-width", 0.5)
    .on("mouseover", function(event, d) {
      const name = d.properties.name;
      const val = valueByCountry[name];
      tooltip.html(`<strong>${name}</strong><br/>${val != null ? val + "%" : "No data"}`)
        .style("left", (event.pageX + 10) + "px")
        .style("top", (event.pageY - 40) + "px")
        .transition().duration(300)
        .style("opacity", 1)
        .style("transform", "translateY(-10px)");
    })
    .on("mouseout", function() {
      tooltip.transition().duration(500)
        .style("opacity", 0)
        .style("transform", "translateY(0px)");
    });

  // Add zoom and pan
  const zoomBehavior = d3.zoom()
    .scaleExtent([1, 8])
    .on("zoom", (event) => {
      g.attr("transform", event.transform);
    });
  svg.call(zoomBehavior);

  // Zoom to continent if selected
  if (selectedContinent && selectedContinent !== "All") {
    // Get country names in the selected continent
    const countriesInContinent = data
      .filter(d => d.CONTINENT === selectedContinent)
      .map(d => d["Reference area"]);
    // Find features for those countries
    const features = countries.filter(f => countriesInContinent.includes(f.properties.name));
    if (features.length > 0) {
      // Compute bounding box
      const geoPath = d3.geoPath().projection(projection);
      const bounds = geoPath.bounds({type: "FeatureCollection", features});
      const [[x0, y0], [x1, y1]] = bounds;
      const dx = x1 - x0, dy = y1 - y0;
      const scale = Math.max(1, Math.min(8, 0.9 / Math.max(dx / width, dy / height)));
      const translate = [
        width / 2 - scale * (x0 + dx / 2),
        height / 2 - scale * (y0 + dy / 2)
      ];
      svg.transition().duration(1000).call(
        zoomBehavior.transform,
        d3.zoomIdentity.translate(translate[0], translate[1]).scale(scale)
      );
    }
  }

  // Add legend for color scale (draw OUTSIDE zoom group, so it remains fixed)
  const legendWidth = 200;
  const legendHeight = 12;
  // Remove any previous choropleth legend
  svg.selectAll(".choropleth-legend").remove();
  // Place legend at the bottom center of the SVG
  const legendSvg = svg.append("g")
    .attr("class", "choropleth-legend")
    .attr("transform", `translate(${(+svg.attr("width") - 200) / 2},${+svg.attr("height") - 40})`);

  const defs = svg.select("defs").empty() ? svg.append("defs") : svg.select("defs");
  defs.selectAll("#choropleth-gradient").remove();
  const linearGradient = defs.append("linearGradient")
    .attr("id", "choropleth-gradient");
  linearGradient.selectAll("stop")
    .data(d3.range(0, 1.01, 0.01))
    .enter().append("stop")
    .attr("offset", d => d)
    .attr("stop-color", d => colorScale(d * 100));

  legendSvg.append("rect")
    .attr("width", legendWidth)
    .attr("height", legendHeight)
    .style("fill", "url(#choropleth-gradient)");

  // Legend axis
  const legendScale = d3.scaleLinear().domain([0, 100]).range([0, legendWidth]);
  const legendAxis = d3.axisBottom(legendScale).ticks(5).tickFormat(d => d + "%");
  legendSvg.append("g")
    .attr("transform", `translate(0,${legendHeight})`)
    .call(legendAxis);

  // Dark mode styling for legend axis and labels
  if (darkMode) {
    legendSvg.selectAll("text").style("fill", "#eee");
    legendSvg.selectAll("path").style("stroke", "#eee");
    legendSvg.selectAll("line").style("stroke", "#eee");
  } else {
    legendSvg.selectAll("text").style("fill", "#000");
    legendSvg.selectAll("path").style("stroke", "#000");
    legendSvg.selectAll("line").style("stroke", "#000");
  }
}