labelimage/components/label/useDescToolsStore.ts

import { Project } from "./api/project/typing"
import { create } from "zustand"
import { splitWord } from "./components/EditorContainer"

interface SubAttribute {
  value?: string | { [x: string]: any }
  is_correct?: boolean
  comment?: string
}

interface QaDataProps {
  value: string
  id: number // 轮次
  is_pre: boolean // 是否是预设问题
  tag: number // 标识该问题是否被审核过 0-未审核 1-正确 2-错误
  uid: number // 创建人
  create_timestamp: number // 创建时间
  modify_uid: number // 修改人
  modify_timestamp: number // 修改时间
  comment: string // 批注内容
  flag: number // 标识被批注过的问题，标注员是否已修改
}

interface DescToolsStore {
  updateDescDataFlag: boolean
  descOperations: Project.LabelSchemaList[]
  setDescOperations: (val: Project.LabelSchemaList[]) => void
  metaOperation: Project.LabelSchemaList | null
  setMetaOperation: (val: Project.LabelSchemaList | null) => void
  descData: Map<string, Map<string, SubAttribute>>
  getFormDataByImageId: (id: string) => { [x: string]: any }
  getDataByClassName: (id: string, name: string) => SubAttribute | null
  setDescData: (data: Map<string, Map<string, SubAttribute>>) => void
  metaData: Map<string, string[]>
  setMetaData: (data: Map<string, string[]>) => void
  qaOperations: Project.LabelSchemaList[]
  setQaOperations: (val: Project.LabelSchemaList[]) => void
  qaData: Map<string, { [x: string]: Array<{ [x: string]: QaDataProps }> }>
  setQaData: (
    data: Map<string, { [x: string]: Array<{ [x: string]: QaDataProps }> }>
  ) => void
  getQuestionList: (imageId: string) => string[]
  countTurnsAndQaNumber: (
    imageId: string,
    label: string
  ) => {
    turns: number
    question_size: number
  }
  countAllTurnsAndQaNumber: (imageId: string) => {
    turns: number
    question_size: number
    new_single_turns: number
    new_multiple_turns: number
  }
  resetData: () => void
  updateFlag: (val: boolean) => void
}

export const useDescToolsStore = create<DescToolsStore>((set) => ({
  updateDescDataFlag: false,
  descOperations: [],
  descData: new Map(),
  metaOperation: null,
  metaData: new Map(),
  qaOperations: [],
  qaData: new Map(),
  setDescOperations: (val) =>
    set((state) => ({
      ...state,
      descOperations: val,
    })),
  getFormDataByImageId: (id) => {
    let obj: any = {}
    const data: Map<string, SubAttribute> | undefined = useDescToolsStore
      .getState()
      .descData.get(id)
    if (!data) return {}
    data.entries().forEach(([label_class, subattr]) => {
      if (typeof subattr.value === "string") {
        obj[label_class] = {
          ...subattr,
          value: subattr.value.split(splitWord)[0],
        }
      } else if (typeof subattr.value === "object") {
        obj[label_class] = { ...subattr.value }
      } else {
        obj[label_class] = {}
      }
    })
    return obj
  },
  getDataByClassName: (id, label_class) => {
    let attr = null
    const data: any = useDescToolsStore.getState().descData.get(id)
    if (data && data.has(label_class)) {
      attr = data.get(label_class)
    }
    return attr
  },
  setDescData: (data) =>
    set((state) => ({
      ...state,
      descData: data,
    })),
  setMetaOperation: (val) =>
    set((state) => ({
      ...state,
      metaOperation: val,
    })),
  setMetaData: (data) =>
    set((state) => ({
      ...state,
      metaData: data,
    })),
  setQaOperations: (val) =>
    set((state) => ({
      ...state,
      qaOperations: val,
    })),
  setQaData: (data) =>
    set((state) => ({
      ...state,
      qaData: data,
    })),
  getQuestionList: (imageId) => {
    let names: string[] = []
    const data = useDescToolsStore.getState().qaData.get(imageId) ?? {}
    Object.values(data).forEach((value) => {
      value.forEach((v) => {
        const key = Object.keys(v)[0]
        names.push(key)
      })
    })
    return names
  },
  countTurnsAndQaNumber: (imageId, label) => {
    let obj = {
      turns: 0,
      question_size: 0,
    }
    const data =
      useDescToolsStore.getState().qaData.get(imageId) ??
      useDescToolsStore.getState().qaData.get("init") ??
      {}
    const labelArr = data[label]
    if (labelArr && labelArr.length) {
      obj.question_size = labelArr.length
      obj.turns = [
        ...new Set(
          labelArr.map((l) => {
            const v = Object.values(l)[0]
            return v.id
          })
        ),
      ].length
    }
    return obj
  },
  countAllTurnsAndQaNumber: (imageId) => {
    let obj = {
      turns: 0,
      question_size: 0,
      new_single_turns: 0,
      new_multiple_turns: 0,
    }
    const data =
      useDescToolsStore.getState().qaData.get(imageId) ??
      useDescToolsStore.getState().qaData.get("init") ??
      {}
    Object.values(data).forEach((value) => {
      obj.question_size += value.length
      obj.turns += [
        ...new Set(
          value.map((d) => {
            const v = Object.values(d)[0]
            return v.id
          })
        ),
      ].length
      let singleRoundCount = 0
      let multiRoundCount = 0
      const roundMap = new Map<
        number,
        { hasNewQuestion: boolean; count: number }
      >()
      value.forEach((d) => {
        const v = Object.values(d)[0]
        if (!roundMap.has(v.id)) {
          roundMap.set(v.id, { hasNewQuestion: !v.is_pre, count: 1 })
        } else {
          const exist = roundMap.get(v.id)
          if (!exist?.hasNewQuestion && !v.is_pre) exist!.hasNewQuestion = true
          exist!.count++
        }
      })
      roundMap.values().forEach((round) => {
        if (round.hasNewQuestion) {
          if (round.count > 1) {
            multiRoundCount++
          } else {
            singleRoundCount++
          }
        }
      })
      obj.new_single_turns += singleRoundCount
      obj.new_multiple_turns += multiRoundCount
    })
    return obj
  },
  resetData: () =>
    set((state) => ({
      ...state,
      descData: new Map(),
      metaData: new Map(),
    })),
  updateFlag: (data) =>
    set((state) => ({
      ...state,
      updateDescDataFlag: data,
    })),
}))