From 2f1cb918862daccbcf2b7b835dfe874f4283f7f8 Mon Sep 17 00:00:00 2001
From: moririn2528 <strangenoise1@gmail.com>
Date: Fri, 8 Sep 2023 16:57:11 +0900
Subject: [PATCH] add undo button on feedback screen

---
 .../ts/components/PreferenceHistory.tsx       |  8 +--
 .../ts/components/PreferentialTrials.tsx      | 49 +++++++++++++++----
 2 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/optuna_dashboard/ts/components/PreferenceHistory.tsx b/optuna_dashboard/ts/components/PreferenceHistory.tsx
index bfacfc62..0535dc2e 100644
--- a/optuna_dashboard/ts/components/PreferenceHistory.tsx
+++ b/optuna_dashboard/ts/components/PreferenceHistory.tsx
@@ -145,16 +145,17 @@ const ChoiceTrials: FC<{
   const [enabled, setEnabled] = useState(choice.enabled)
   const theme = useTheme()
   const worst_trials = new Set([choice.clicked])
-  const actions = actionCreator()
+  const action = actionCreator()
   const handleSwitch = () => {
     setEnabled(!enabled)
-    actions.switchPreferentialHistory(study_id, choice.id, !enabled)
+    action.switchPreferentialHistory(study_id, choice.id, !enabled)
   }
 
   return (
     <Box
       sx={{
         marginBottom: theme.spacing(4),
+        position: "relative",
       }}
     >
       <Box
@@ -180,7 +181,7 @@ const ChoiceTrials: FC<{
             margin: `auto ${theme.spacing(2)}`,
           }}
         >
-          {enabled ? <DeleteIcon /> : <RestoreFromTrashIcon />}
+          {choice.enabled ? <DeleteIcon /> : <RestoreFromTrashIcon />}
         </IconButton>
       </Box>
       <Box
@@ -188,6 +189,7 @@ const ChoiceTrials: FC<{
           display: "flex",
           flexDirection: "row",
           flexWrap: "wrap",
+          filter: choice.enabled ? undefined : "brightness(0.4)",
         }}
       >
         {choice.candidates.map((trial_num, index) => (
diff --git a/optuna_dashboard/ts/components/PreferentialTrials.tsx b/optuna_dashboard/ts/components/PreferentialTrials.tsx
index 08efd9a7..ab7f04ea 100644
--- a/optuna_dashboard/ts/components/PreferentialTrials.tsx
+++ b/optuna_dashboard/ts/components/PreferentialTrials.tsx
@@ -14,6 +14,7 @@ import OpenInFullIcon from "@mui/icons-material/OpenInFull"
 import ReplayIcon from "@mui/icons-material/Replay"
 import Modal from "@mui/material/Modal"
 import { red } from "@mui/material/colors"
+import UndoIcon from "@mui/icons-material/Undo"
 
 import { actionCreator } from "../action"
 import { TrialListDetail } from "./TrialList"
@@ -180,10 +181,13 @@ type DisplayTrials = {
 export const PreferentialTrials: FC<{ studyDetail: StudyDetail | null }> = ({
   studyDetail,
 }) => {
+  const [undoHistoryId, setUndoHistoryId] = useState<string | null>(null)
+
   if (studyDetail === null || !studyDetail.is_preferential) {
     return null
   }
   const theme = useTheme()
+  const action = actionCreator()
 
   const runningTrials = studyDetail.trials.filter((t) => t.state === "Running")
   const activeTrials = runningTrials.concat(studyDetail.best_trials)
@@ -229,18 +233,45 @@ export const PreferentialTrials: FC<{ studyDetail: StudyDetail | null }> = ({
       }
     })
   }
+  const latestHistoryId = studyDetail?.preference_history
+    ?.filter((h) => h.enabled)
+    .pop()?.id
+  if (undoHistoryId !== null && undoHistoryId !== latestHistoryId) {
+    setUndoHistoryId(null)
+  }
 
   return (
     <Box padding={theme.spacing(2)}>
-      <Typography
-        variant="h4"
-        sx={{
-          marginBottom: theme.spacing(2),
-          fontWeight: theme.typography.fontWeightBold,
-        }}
-      >
-        Which trial is the worst?
-      </Typography>
+      <Box display="flex">
+        <Typography
+          variant="h4"
+          sx={{
+            marginBottom: theme.spacing(2),
+            fontWeight: theme.typography.fontWeightBold,
+          }}
+        >
+          Which trial is the worst?
+        </Typography>
+        <IconButton
+          disabled={latestHistoryId === undefined || undoHistoryId !== null}
+          onClick={() => {
+            if (latestHistoryId === undefined) {
+              return
+            }
+            setUndoHistoryId(latestHistoryId)
+            action.switchPreferentialHistory(
+              studyDetail.id,
+              latestHistoryId,
+              false
+            )
+          }}
+          sx={{
+            margin: "auto 0 auto auto",
+          }}
+        >
+          <UndoIcon />
+        </IconButton>
+      </Box>
       <Box sx={{ display: "flex", flexDirection: "row", flexWrap: "wrap" }}>
         {displayTrials.numbers.map((t, index) => (
           <PreferentialTrial