From f5a3db26b0078f4a1b8295fe2e412d6f9167b49d Mon Sep 17 00:00:00 2001
From: jorgenherje <jorgen.herje@ceetronsolutions.com>
Date: Wed, 6 Sep 2023 14:11:02 +0200
Subject: [PATCH] Fixup based on feedback after testing

- Added ApiStatesWrapper for list of react queries
- Handle selected vectors when deleting/removing ensembles
- Handle selected vectors when selecting new or non-cached ensemble
- Fix bug in use effect in settings, was `return` instead of `continue` in for-loop, thus not triggering update of view.
---
 .../ApiStatesWrapper/apiStatesWrapper.tsx     | 41 +++++++++++
 .../lib/components/ApiStatesWrapper/index.ts  |  1 +
 .../SimulationTimeSeriesMatrix/settings.tsx   | 69 +++++++++++--------
 .../SimulationTimeSeriesMatrix/view.tsx       |  9 ++-
 4 files changed, 90 insertions(+), 30 deletions(-)
 create mode 100644 frontend/src/lib/components/ApiStatesWrapper/apiStatesWrapper.tsx
 create mode 100644 frontend/src/lib/components/ApiStatesWrapper/index.ts

diff --git a/frontend/src/lib/components/ApiStatesWrapper/apiStatesWrapper.tsx b/frontend/src/lib/components/ApiStatesWrapper/apiStatesWrapper.tsx
new file mode 100644
index 000000000..7d88d7898
--- /dev/null
+++ b/frontend/src/lib/components/ApiStatesWrapper/apiStatesWrapper.tsx
@@ -0,0 +1,41 @@
+import React from "react";
+
+import { resolveClassNames } from "@lib/utils/resolveClassNames";
+import { QueryObserverResult } from "@tanstack/react-query";
+
+export type ApiStatesWrapperProps = {
+    apiResults: QueryObserverResult[];
+    loadingComponent: React.ReactNode;
+    errorComponent: React.ReactNode;
+    className?: string;
+    style?: React.CSSProperties;
+    children: React.ReactNode;
+};
+
+export const ApiStatesWrapper: React.FC<ApiStatesWrapperProps> = (props: ApiStatesWrapperProps) => {
+    return (
+        <div
+            className={resolveClassNames(
+                "relative rounded",
+                { "outline outline-blue-100 outline-offset-2": props.apiResults.some((elm) => elm.isLoading) },
+                { "outline outline-red-100 outline-offset-2": props.apiResults.some((elm) => elm.isError) },
+                props.className ?? ""
+            )}
+            style={props.style}
+        >
+            {props.apiResults.some((elm) => elm.isLoading) && (
+                <div className="absolute left-0 right-0 w-full h-full bg-white bg-opacity-80 flex items-center justify-center z-10">
+                    {props.loadingComponent}
+                </div>
+            )}
+            {props.apiResults.some((elm) => elm.isError) && (
+                <div className="absolute left-0 right-0 w-full h-full bg-white bg-opacity-80 flex items-center justify-center z-10">
+                    {props.errorComponent}
+                </div>
+            )}
+            {props.children}
+        </div>
+    );
+};
+
+ApiStatesWrapper.displayName = "ApiStatesWrapper";
diff --git a/frontend/src/lib/components/ApiStatesWrapper/index.ts b/frontend/src/lib/components/ApiStatesWrapper/index.ts
new file mode 100644
index 000000000..f81343045
--- /dev/null
+++ b/frontend/src/lib/components/ApiStatesWrapper/index.ts
@@ -0,0 +1 @@
+export { ApiStatesWrapper } from "./apiStatesWrapper";
diff --git a/frontend/src/modules/SimulationTimeSeriesMatrix/settings.tsx b/frontend/src/modules/SimulationTimeSeriesMatrix/settings.tsx
index 3f1e9bcac..09ea1d165 100644
--- a/frontend/src/modules/SimulationTimeSeriesMatrix/settings.tsx
+++ b/frontend/src/modules/SimulationTimeSeriesMatrix/settings.tsx
@@ -6,12 +6,15 @@ import { EnsembleSet } from "@framework/EnsembleSet";
 import { ModuleFCProps } from "@framework/Module";
 import { useEnsembleSet } from "@framework/WorkbenchSession";
 import { MultiEnsembleSelect } from "@framework/components/MultiEnsembleSelect";
+import { fixupEnsembleIdents } from "@framework/utils/ensembleUiHelpers";
+import { ApiStatesWrapper } from "@lib/components/ApiStatesWrapper";
 import { Checkbox } from "@lib/components/Checkbox";
+import { CircularProgress } from "@lib/components/CircularProgress";
 import { CollapsibleGroup } from "@lib/components/CollapsibleGroup";
 import { Dropdown } from "@lib/components/Dropdown";
 import { Label } from "@lib/components/Label";
 import { RadioGroup } from "@lib/components/RadioGroup";
-import { SmartNodeSelectorSelection } from "@lib/components/SmartNodeSelector";
+import { SmartNodeSelectorSelection, TreeDataNode } from "@lib/components/SmartNodeSelector";
 import { VectorSelector } from "@lib/components/VectorSelector";
 import { resolveClassNames } from "@lib/utils/resolveClassNames";
 import { createVectorSelectorDataFromVectors } from "@lib/utils/vectorSelectorUtils";
@@ -46,6 +49,7 @@ export function settings({ moduleContext, workbenchSession }: ModuleFCProps<Stat
     const [previousEnsembleSet, setPreviousEnsembleSet] = React.useState<EnsembleSet>(ensembleSet);
     const [selectedEnsembleIdents, setSelectedEnsembleIdents] = React.useState<EnsembleIdent[]>([]);
     const [selectedVectorNames, setSelectedVectorNames] = React.useState<string[]>([]);
+    const [vectorSelectorData, setVectorSelectorData] = React.useState<TreeDataNode[]>([]);
 
     const [prevVisualizationMode, setPrevVisualizationMode] = React.useState<VisualizationMode>(visualizationMode);
     if (prevVisualizationMode !== visualizationMode) {
@@ -56,27 +60,24 @@ export function settings({ moduleContext, workbenchSession }: ModuleFCProps<Stat
     const ensembleVectorListsHelper = new EnsembleVectorListsHelper(selectedEnsembleIdents, vectorListQueries);
     const vectorsUnion: VectorDescription_api[] = ensembleVectorListsHelper.vectorsUnion();
 
-    const vectorSelectorData = createVectorSelectorDataFromVectors(vectorsUnion.map((vector) => vector.name));
+    const selectedVectorNamesHasHistorical = ensembleVectorListsHelper.hasAnyHistoricalVector(selectedVectorNames);
+    const currentVectorSelectorData = createVectorSelectorDataFromVectors(vectorsUnion.map((vector) => vector.name));
 
-    const newSelectedVectorNames = [];
-    for (const vector of selectedVectorNames) {
-        if (
-            vectorsUnion.some((item) => {
-                return item.name === vector;
-            })
-        ) {
-            newSelectedVectorNames.push(vector);
-        }
-    }
-    if (!isEqual(selectedVectorNames, newSelectedVectorNames)) {
-        setSelectedVectorNames(newSelectedVectorNames);
+    // Only update if all vector lists are retrieved before updating vectorSelectorData has changed
+    const hasVectorListQueriesErrorOrLoading = vectorListQueries.some((query) => query.isLoading || query.isError);
+    if (!hasVectorListQueriesErrorOrLoading && !isEqual(currentVectorSelectorData, vectorSelectorData)) {
+        setVectorSelectorData(currentVectorSelectorData);
     }
 
-    const selectedVectorNamesHasHistorical = ensembleVectorListsHelper.hasAnyHistoricalVector(newSelectedVectorNames);
-
     if (!isEqual(ensembleSet, previousEnsembleSet)) {
-        // TODO:
-        // Handle change of ensembleSet-> validity of ensemble selection and vector selection
+        const newSelectedEnsembleIdents = selectedEnsembleIdents.filter(
+            (ensemble) => ensembleSet.findEnsemble(ensemble) !== null
+        );
+        const validatedEnsembleIdents = fixupEnsembleIdents(newSelectedEnsembleIdents, ensembleSet) ?? [];
+        if (!isEqual(selectedEnsembleIdents, validatedEnsembleIdents)) {
+            setSelectedEnsembleIdents(validatedEnsembleIdents);
+        }
+
         setPreviousEnsembleSet(ensembleSet);
     }
 
@@ -86,7 +87,7 @@ export function settings({ moduleContext, workbenchSession }: ModuleFCProps<Stat
             for (const ensemble of selectedEnsembleIdents) {
                 for (const vector of selectedVectorNames) {
                     if (!ensembleVectorListsHelper.isVectorInEnsemble(ensemble, vector)) {
-                        return;
+                        continue;
                     }
 
                     newVectorSpecifications.push({
@@ -221,15 +222,27 @@ export function settings({ moduleContext, workbenchSession }: ModuleFCProps<Stat
                     disabled={true}
                     onChange={handleShowObservations}
                 />
-                <VectorSelector
-                    data={vectorSelectorData}
-                    selectedTags={selectedVectorNames}
-                    placeholder="Add new vector..."
-                    maxNumSelectedNodes={50}
-                    numSecondsUntilSuggestionsAreShown={0.5}
-                    lineBreakAfterTag={true}
-                    onChange={handleVectorSelectChange}
-                />
+                <div
+                    className={resolveClassNames({
+                        "pointer-events-none opacity-80": vectorListQueries.some((query) => query.isLoading),
+                    })}
+                >
+                    <ApiStatesWrapper
+                        apiResults={vectorListQueries}
+                        loadingComponent={<CircularProgress />}
+                        errorComponent={"Could not load the vectors for selected ensembles"}
+                    >
+                        <VectorSelector
+                            data={vectorSelectorData}
+                            selectedTags={selectedVectorNames}
+                            placeholder="Add new vector..."
+                            maxNumSelectedNodes={50}
+                            numSecondsUntilSuggestionsAreShown={0.5}
+                            lineBreakAfterTag={true}
+                            onChange={handleVectorSelectChange}
+                        />
+                    </ApiStatesWrapper>
+                </div>
             </CollapsibleGroup>
             <CollapsibleGroup expanded={true} title="Visualization">
                 <RadioGroup
diff --git a/frontend/src/modules/SimulationTimeSeriesMatrix/view.tsx b/frontend/src/modules/SimulationTimeSeriesMatrix/view.tsx
index 095b1e323..0b4bc4be9 100644
--- a/frontend/src/modules/SimulationTimeSeriesMatrix/view.tsx
+++ b/frontend/src/modules/SimulationTimeSeriesMatrix/view.tsx
@@ -141,18 +141,23 @@ export const view = ({ moduleContext, workbenchSettings }: ModuleFCProps<State>)
         return;
     }
 
+    const plotData = subplotBuilder.createPlotData();
     // TODO: Keep uirevision?
     return (
         <div className="w-full h-full" ref={wrapperDivRef}>
             <Plot
-                data={subplotBuilder.createPlotData()}
+                key={plotData.length} // Note: To trigger re-render and remove legends when plotData is empty
+                data={plotData}
                 layout={subplotBuilder.createPlotLayout()}
                 config={{ scrollZoom: true }}
                 onHover={handleHover}
                 onUnhover={handleUnHover}
             />
             {isDevMode() && (
-                <div className="absolute top-10 left-5 italic text-pink-400">(rc={renderCount.current})</div>
+                <>
+                    <div className="absolute top-10 left-5 italic text-pink-400">(rc={renderCount.current})</div>
+                    <div className="absolute top-10 left-20 italic text-pink-400"> Traces: {plotData.length}</div>
+                </>
             )}
         </div>
     );