From e65386511501de4d86053d55ea89a87001d18b40 Mon Sep 17 00:00:00 2001
From: "Dae Kun (DK) Kwon" <dkwon@nd.edu>
Date: Tue, 2 Jul 2024 21:30:14 -0400
Subject: [PATCH] export both plotly plot and legend images

---
 .../components/plotControls/PlotLegend.tsx    |  2 +
 .../plotControls/PlotListLegend.tsx           |  2 +
 .../src/plots/ExportPlotToImageButton.tsx     | 40 +++++++++++++++++--
 3 files changed, 41 insertions(+), 3 deletions(-)
 mode change 100644 => 100755 packages/libs/components/src/plots/ExportPlotToImageButton.tsx

diff --git a/packages/libs/components/src/components/plotControls/PlotLegend.tsx b/packages/libs/components/src/components/plotControls/PlotLegend.tsx
index 3391ba561f..ac30853361 100755
--- a/packages/libs/components/src/components/plotControls/PlotLegend.tsx
+++ b/packages/libs/components/src/components/plotControls/PlotLegend.tsx
@@ -34,6 +34,8 @@ export default function PlotLegend({
         type === 'colorscale' ||
         type === 'bubble') && (
         <div
+          // set id for export legend
+          id="plotLegend"
           style={{
             display: 'inline-block', // for general usage (e.g., story)
             border: '1px solid #dedede',
diff --git a/packages/libs/components/src/components/plotControls/PlotListLegend.tsx b/packages/libs/components/src/components/plotControls/PlotListLegend.tsx
index 7f5d0c8ebb..f6065538c5 100644
--- a/packages/libs/components/src/components/plotControls/PlotListLegend.tsx
+++ b/packages/libs/components/src/components/plotControls/PlotListLegend.tsx
@@ -269,6 +269,8 @@ export default function PlotListLegend({
                         whiteSpace: 'nowrap',
                         overflow: 'hidden',
                         textOverflow: 'ellipsis',
+                        // set width to avoid unnecessary ellipsis for image export
+                        width: '100%',
                       }}
                     >
                       {item.label === 'No data' ||
diff --git a/packages/libs/components/src/plots/ExportPlotToImageButton.tsx b/packages/libs/components/src/plots/ExportPlotToImageButton.tsx
old mode 100644
new mode 100755
index cc022c195a..be7a40372c
--- a/packages/libs/components/src/plots/ExportPlotToImageButton.tsx
+++ b/packages/libs/components/src/plots/ExportPlotToImageButton.tsx
@@ -1,6 +1,7 @@
 import { Image } from '@material-ui/icons';
 import { colors, SingleSelect, Warning } from '@veupathdb/coreui';
 import { CSSProperties, useState } from 'react';
+import domtoimage from 'dom-to-image';
 
 interface ToImageOpts {
   height: number;
@@ -63,11 +64,34 @@ export function ExportPlotToImageButton(props: Props) {
           if (format) {
             setSawError(false);
             try {
-              await downloadImage(toImage, filename, {
+              const options = {
                 format,
                 height: imageHeight,
                 width: imageWidth,
-              });
+              };
+              const legendStyles = {
+                style: {
+                  width: '400px',
+                  backgroundColor: '#fff',
+                },
+              };
+
+              // export plotly plot
+              await downloadImage(toImage, filename, options);
+
+              // export legend
+              if (format === 'svg')
+                await domtoimage
+                  .toSvg(document.getElementById('plotLegend')!, legendStyles)
+                  .then((dataUrl) => {
+                    saveImage(dataUrl, filename, options, true);
+                  });
+              else if (format === 'png')
+                await domtoimage
+                  .toPng(document.getElementById('plotLegend')!, legendStyles)
+                  .then((dataUrl) => {
+                    saveImage(dataUrl, filename, options, true);
+                  });
             } catch (error) {
               setSawError(true);
               console.error(error);
@@ -102,9 +126,19 @@ async function downloadImage(
   options: ToImageOpts
 ) {
   const imgUrl = await toImage(options);
+  saveImage(imgUrl, filename, options);
+}
+
+function saveImage(
+  imgUrl: string,
+  filename: string,
+  options: ToImageOpts,
+  isLegend: boolean = false
+) {
   const downloadLink = document.createElement('a');
   downloadLink.href = imgUrl;
-  downloadLink.download = filename + '.' + options.format;
+  downloadLink.download =
+    filename + (isLegend ? 'Legend' : '') + '.' + options.format;
   document.body.appendChild(downloadLink);
   downloadLink.click();
   document.body.removeChild(downloadLink);