All files / src/pages/flowsheet-page/flowsheet/PropertiesSidebar/MachineLearning MachineLearningData.tsx

92.3% Statements 12/13
84.61% Branches 11/13
0% Functions 0/1
91.66% Lines 11/12

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86                                            31x   31x 31x 31x       31x   27x   27x 27x   27x                                       27x             27x                                          
import { useState } from "react";
import { Button } from "@/ahuora-design-system/ui/button";
import {
  Dialog,
  DialogContent,
  DialogTrigger,
} from "@/ahuora-design-system/ui/dialog";
import { useCoreMlListQuery } from "@/api/apiStore.gen";
import { useSearchParam } from "@/hooks/searchParams";
import { cn } from "@/lib/utils";
import ColumnMappings from "./ColumnMappings";
import ProgressBar from "./ProgressBar";
import SuccesslyTrained from "./SuccesslyTrained";
import Train from "./Train";
import TrainingFailed from "./TrainingFailed";
import UploadCSV, { type CsvHeaderPreviewState } from "./UploadCSV";
 
type MlModelWithCsvMeta = {
  csv_file_name?: string;
};
 
export default function MachineLearningData() {
  const [objectId] = useSearchParam("object");
  const [csvHeaderPreview, setCsvHeaderPreview] =
    useState<CsvHeaderPreviewState | null>(null);
  const simulationObject = +(objectId || 0);
  const { data: model } = useCoreMlListQuery({
    simulationObject: simulationObject,
  });
 
  if (!model) return null;
 
  const currentModel = model?.[0] as (typeof model)[number] &
    MlModelWithCsvMeta;
  const progress = currentModel?.progress ?? 0;
  const modelId = currentModel?.id ?? 0;
 
  const contentByProgress: Record<number, React.ReactNode> = {
    0: (
      <UploadCSV
        simulationObject={simulationObject}
        onCsvHeaderPreviewChange={setCsvHeaderPreview}
      />
    ),
    1: (
      <ColumnMappings
        modelId={modelId}
        simulationObject={simulationObject}
        csvFileName={currentModel?.csv_file_name}
        csvHeaderPreview={csvHeaderPreview}
      />
    ),
    2: <Train />,
    3: <SuccesslyTrained model={currentModel} />,
    4: <TrainingFailed />,
  };
 
  const dialogContent = contentByProgress[progress] ?? (
    <UploadCSV
      simulationObject={simulationObject}
      onCsvHeaderPreviewChange={setCsvHeaderPreview}
    />
  );
 
  return (
    <Dialog>
      <div className="w-full flex justify-center p-4">
        <DialogTrigger asChild>
          <Button>{progress !== 0 ? "View" : "Set-up"}</Button>
        </DialogTrigger>
      </div>
      <DialogContent
        preventDefault
        className={cn(
          "flex w-[calc(100vw-2rem)] max-w-[880px] flex-col",
          progress === 3 ? "h-[78vh] max-h-[860px]" : "h-[65%]",
        )}
        onInteractOutside={(event) => event.preventDefault()}
      >
        <ProgressBar progress={progress} />
        {dialogContent}
      </DialogContent>
    </Dialog>
  );
}