From 2086bbd14c755a60544b86276d94db39f9a4327b Mon Sep 17 00:00:00 2001
From: c-bata <contact@c-bata.link>
Date: Thu, 7 Sep 2023 18:59:10 +0900
Subject: [PATCH] Simplify ThreeJsArtifactViewer

---
 .../ts/components/ThreejsArtifactViewer.tsx   | 48 ++++++++++++++++-
 optuna_dashboard/ts/components/TrialList.tsx  | 52 ++++---------------
 2 files changed, 56 insertions(+), 44 deletions(-)

diff --git a/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx b/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx
index 536f1d91..ec30e63f 100644
--- a/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx
+++ b/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx
@@ -1,10 +1,11 @@
 import * as THREE from "three"
-import React, { useEffect, useState } from "react"
+import React, { useEffect, useState, ReactNode } from "react"
 import { Canvas } from "@react-three/fiber"
 import { GizmoHelper, GizmoViewport, OrbitControls } from "@react-three/drei"
 import { STLLoader } from "three/examples/jsm/loaders/STLLoader"
 import { Rhino3dmLoader } from "three/examples/jsm/loaders/3DMLoader"
 import { PerspectiveCamera } from "three"
+import { Modal, Box } from "@mui/material"
 
 interface ThreejsArtifactViewerProps {
   src: string
@@ -109,3 +110,48 @@ export const ThreejsArtifactViewer: React.FC<ThreejsArtifactViewerProps> = (
     </Canvas>
   )
 }
+
+export const useThreejsArtifactModal = (): [
+  (path: string, artifact: Artifact) => void,
+  () => ReactNode
+] => {
+  const [open, setOpen] = useState(false)
+  const [target, setTarget] = useState<[string, Artifact | null]>(["", null])
+
+  const openModal = (artifactUrlPath: string, artifact: Artifact) => {
+    setTarget([artifactUrlPath, artifact])
+    setOpen(true)
+  }
+
+  const renderDeleteStudyDialog = () => {
+    return (
+      <Modal
+        open={open}
+        onClose={() => {
+          setOpen(false)
+          setTarget(["", null])
+        }}
+      >
+        <Box
+          sx={{
+            position: "absolute",
+            top: "50%",
+            left: "50%",
+            transform: "translate(-50%, -50%)",
+            bgcolor: "background.paper",
+            borderRadius: "15px",
+          }}
+        >
+          <ThreejsArtifactViewer
+            src={target[0]}
+            width={`${innerWidth * 0.8}px`}
+            height={`${innerHeight * 0.8}px`}
+            hasGizmo={true}
+            filetype={target[1]?.filename.split(".").pop()}
+          />
+        </Box>
+      </Modal>
+    )
+  }
+  return [openModal, renderDeleteStudyDialog]
+}
diff --git a/optuna_dashboard/ts/components/TrialList.tsx b/optuna_dashboard/ts/components/TrialList.tsx
index 4aceb619..84f98c0c 100644
--- a/optuna_dashboard/ts/components/TrialList.tsx
+++ b/optuna_dashboard/ts/components/TrialList.tsx
@@ -20,7 +20,6 @@ import {
   CardContent,
   CardMedia,
   CardActionArea,
-  Modal,
 } from "@mui/material"
 import Chip from "@mui/material/Chip"
 import Divider from "@mui/material/Divider"
@@ -47,7 +46,10 @@ import { artifactIsAvailable } from "../state"
 import { actionCreator } from "../action"
 import { useDeleteArtifactDialog } from "./DeleteArtifactDialog"
 import { TrialFormWidgets } from "./TrialFormWidgets"
-import { ThreejsArtifactViewer } from "./ThreejsArtifactViewer"
+import {
+  ThreejsArtifactViewer,
+  useThreejsArtifactModal,
+} from "./ThreejsArtifactViewer"
 
 const states: TrialState[] = [
   "Complete",
@@ -330,9 +332,8 @@ const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => {
   const [openDeleteArtifactDialog, renderDeleteArtifactDialog] =
     useDeleteArtifactDialog()
   const [dragOver, setDragOver] = useState<boolean>(false)
-  const [open3dModelViewer, setOpen3dModelViewer] = useState<{
-    [key: string]: boolean
-  }>({})
+  const [openThreejsArtifactModal, renderThreejsArtifactModal] =
+    useThreejsArtifactModal()
 
   const width = "200px"
   const height = "150px"
@@ -499,48 +500,12 @@ const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => {
                     color="inherit"
                     sx={{ margin: "auto 0" }}
                     onClick={() => {
-                      setOpen3dModelViewer(() => {
-                        const obj = { ...open3dModelViewer }
-                        obj[a.artifact_id] = true
-                        return obj
-                      })
+                      const urlPath = `/artifacts/${trial.study_id}/${trial.trial_id}/${a.artifact_id}`
+                      openThreejsArtifactModal(urlPath, a)
                     }}
                   >
                     <FullscreenIcon />
                   </IconButton>
-                  <Modal
-                    open={
-                      a.artifact_id in open3dModelViewer
-                        ? open3dModelViewer[a.artifact_id]
-                        : false
-                    }
-                    onClose={() => {
-                      setOpen3dModelViewer(() => {
-                        const obj = { ...open3dModelViewer }
-                        obj[a.artifact_id] = false
-                        return obj
-                      })
-                    }}
-                  >
-                    <Box
-                      sx={{
-                        position: "absolute",
-                        top: "50%",
-                        left: "50%",
-                        transform: "translate(-50%, -50%)",
-                        bgcolor: "background.paper",
-                        borderRadius: "15px",
-                      }}
-                    >
-                      <ThreejsArtifactViewer
-                        src={`/artifacts/${trial.study_id}/${trial.trial_id}/${a.artifact_id}`}
-                        width={`${innerWidth * 0.8}px`}
-                        height={`${innerHeight * 0.8}px`}
-                        hasGizmo={true}
-                        filetype={a.filename.split(".").pop()}
-                      />
-                    </Box>
-                  </Modal>
                   <IconButton
                     aria-label="delete artifact"
                     size="small"
@@ -763,6 +728,7 @@ const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => {
         ) : null}
       </Box>
       {renderDeleteArtifactDialog()}
+      {renderThreejsArtifactModal()}
     </>
   )
 }