import { describe, expect, it } from "vitest"
import { batchSizeFromTensorShape, canonicalizeBatchTensorShape } from "../../web/logic/any-split-batch.js"
import {
  buildSplitOutputValues,
  computeDesiredDataOutputCount,
  computeSplitNumData,
} from "../../web/logic/any-split-layout.js"

describe("any_split", () => {
  describe("canonicalizeBatchTensorShape", () => {
    it("[B,1,H,W] の MASK バッチを [B,H,W] に正規化", () => {
      expect(canonicalizeBatchTensorShape(4, [3, 1, 993, 1000])).toEqual({
        ndim: 3,
        shape: [3, 993, 1000],
      })
    })

    it("IMAGE [B,H,W,C] は変更しない", () => {
      expect(canonicalizeBatchTensorShape(4, [3, 512, 512, 3])).toEqual({
        ndim: 4,
        shape: [3, 512, 512, 3],
      })
    })

    it("2D MASK [H,W] は変更しない", () => {
      expect(canonicalizeBatchTensorShape(2, [993, 1000])).toEqual({
        ndim: 2,
        shape: [993, 1000],
      })
    })
  })

  describe("batchSizeFromTensorShape", () => {
    it("[B,1,H,W] は B を返す", () => {
      expect(batchSizeFromTensorShape(4, [3, 1, 993, 1000])).toBe(3)
    })

    it("2D MASK は 1", () => {
      expect(batchSizeFromTensorShape(2, [993, 1000])).toBe(1)
    })
  })

  describe("computeSplitNumData", () => {
    it("接続なしのとき any_01 のみ", () => {
      expect(computeSplitNumData([])).toBe(1)
    })

    it("length のみ接続のとき data 本数は増えない", () => {
      expect(computeSplitNumData([1])).toBe(1)
    })

    it("any_01 接続で any_02 まで確保", () => {
      expect(computeSplitNumData([0])).toBe(2)
    })

    it("any_03 接続で any_04 まで確保", () => {
      expect(computeSplitNumData([2])).toBe(4)
    })

    it("any_01 と length 接続でも data 本数は 2", () => {
      expect(computeSplitNumData([0, 1])).toBe(2)
    })
  })

  describe("computeDesiredDataOutputCount", () => {
    it("data 出力が未接続なら 1", () => {
      expect(computeDesiredDataOutputCount([{}])).toBe(1)
    })

    it("any_01 接続で 2", () => {
      expect(computeDesiredDataOutputCount([{ links: [1] }])).toBe(2)
    })

    it("any_03 接続で 4", () => {
      const outputs = [{}, {}, { links: [1] }]
      expect(computeDesiredDataOutputCount(outputs)).toBe(4)
    })
  })

  describe("buildSplitOutputValues", () => {
    it("3 要素リストを分解し length=3", () => {
      const items = ["a", "b", "c"]
      const { values, length } = buildSplitOutputValues(items, [0])
      expect(length).toBe(3)
      expect(values[0]).toBe("a")
      expect(values[1]).toBe("b")
      expect(values[2]).toBe(3)
    })

    it("空リストのとき length のみ", () => {
      const { values, length } = buildSplitOutputValues([], [])
      expect(length).toBe(0)
      expect(values).toEqual([0])
    })
  })
})