import { _detectObjects, _tagImage, _captionImage, _embedImage, _findRegions, _visionFile } from "agency-lang/stdlib-lib/vision.js"

/** @module
  @summary Detect, tag, caption, and compare images with local vision models.
  Ask a local vision model about an image: which objects are where, which
booru tags describe it, a caption, or embeddings that say how parts of
it look. The model must be running: start it with
`agency local serve <model>`. The image never leaves the machine.

  ```ts
  import { detectObjects, tagImage } from "std::vision"
  import { cropImage } from "std::image"

  node main() {
    handle {
      const found = detectObjects("page.png", ["person", "desk", "chair"], "florence-2") catch []
      for (hit in found) {
        cropImage("page.png", hit.box, "crops/${hit.label}_${hit.id}.png", pad: 0.05)
      }
      const tags = tagImage("crops/person_0.png", "wd14-tagger") catch []
      const names = map(tags) as t {
        return t.tag
      }
      print(names.join(", "))
    } with approve
  }
  ```

Each function is one job with a `Result`, so a labeling loop is a few
lines of your own code over them, and partial application narrows each:
`detectObjects.partial(labels: ["person"], model: "florence-2")` is a
person-finder.

## Finding your own things

A detector finds things by the name of their kind. It finds every mug,
and cannot tell yours from the one next to it. To find your own mug, or
the cats you draw, compare instead:

1. `findRegions` boxes every thing in a picture, with no names. Where a
   detector knows the kind, `detectObjects` gives the boxes instead.
2. `embedImage` turns each box into an embedding: a list of numbers
   that says how it looks. Things that look alike get embeddings that
   are close together.
3. Compare each box with embeddings of crops of your own things, and of
   other things, and keep the boxes nearest yours.

```ts
import { findRegions, embedImage } from "std::vision"
import { cosineSimilarity } from "std::embedding"

// How alike `vector` is to the closest of `references`, where 1 is identical.
def closest(vector: number[], references: number[][]): number {
  let best = 0
  for (reference in references) {
    const similarity = cosineSimilarity(vector, reference)
    if (isSuccess(similarity) && similarity.value > best) {
      best = similarity.value
    }
  }
  return best
}

// One vector per crop file.
def embedCrops(paths: string[]): number[][] {
  let vectors: number[][] = []
  for (path in paths) {
    const embedded = embedImage(path, "dinov2-base")
    if (isSuccess(embedded)) {
      vectors.push(embedded.value[0])
    }
  }
  return vectors
}

node main(page: string, catCrops: string[], otherCrops: string[]) {
  const cats = embedCrops(catCrops)
  const others = embedCrops(otherCrops)
  const regions = findRegions(page, "owlv2-base")
  if (isFailure(regions)) {
    return regions.error
  }
  const boxes = map(regions.value) as region {
    return region.box
  }
  const vectors = embedImage(page, "dinov2-base", boxes: boxes)
  if (isFailure(vectors)) {
    return vectors.error
  }
  for (vector, i in vectors.value) {
    if (closest(vector, cats) > closest(vector, others)) {
      print("Cat in region ${i}")
    }
  }
}
```

The crops are tight crops made with `cropImage`, one file each. A crop
file embedded whole is cut the same way as a box embedded from a page,
so the two compare fairly. A whole photo of a mug on a desk and a tight
crop of a mug do not: the first vector mostly describes the desk.
*/

// One effect for the five functions, because all five do the same thing
// to the file: read it once and hand it to a local model. The task and
// the model are in the payload for a policy that wants them.
@alwaysUnder(dir)
effect std::vision { dir: string, filename: string, task: string, model: string }

/** Normalized to 0..1 with the origin at the top left of the image, so
`y` grows downward. The same shape std::ocr returns, so a box from
either can go to cropImage. */
export type BoundingBox = {
  x: number,
  y: number,
  width: number,
  height: number,
}

/** One thing a detector found. `id` is its index in the reply, so crops
can be named after it. The box is normalized 0..1 from the top left. */
export type Detection = {
  id: number,
  label: string,
  score: number,
  box: BoundingBox,
}

/** One box around a thing in a picture, found with no name. `id` is its
index in the reply, so crops can be named after it. The score says how
much it looks like a thing at all, 0..1. */
export type Region = {
  id: number,
  score: number,
  box: BoundingBox,
}

/** One tag a tagger gave, with how sure it was, 0..1. */
export type Tag = {
  tag: string,
  score: number,
}

/**
Use labels for object categories such as "person" or "chair". OWLv2 handles
the labels in one pass. Florence-2 makes a separate pass for each label.
Florence-2 assigns every detection a score of 1. OWLv2 uses a default
threshold of 0.1 when none is supplied.
*/
@capabilities(fileRead, network)
export idempotent def detectObjects(
  path: string,
  labels: string[],
  model: string,
  threshold: number | null = null,
): Result<Detection[]> raises <std::vision> {
  """
  Locate named objects in an image using a local detector.

  @param path - Path to a PNG, JPEG, WebP, or GIF image.
  @param labels - One to 50 object names to look for, such as "person" or "chair".
  @param model - A running local detector, such as "owlv2-base" or "florence-2".
  @param threshold - Minimum score from 0 to 1, or null for the model default.
  """
  const file = try _visionFile(path)
  if (isFailure(file)) {
    return file
  }
  return interrupt std::vision("Look for ${labels.join(", ")} in this image with a local model?", {
    dir: file.value.dir,
    filename: file.value.filename,
    task: "detections",
    model: model,
  })
  return try _detectObjects(file.value.image, labels, model, threshold)
}

/**
Use the returned booru tags in illustration-model prompts or caption files.
Tags arrive in descending score order.
*/
@capabilities(fileRead, network)
export idempotent def tagImage(
  path: string,
  model: string,
  threshold: number = 0.35,
  limit: number = 30,
): Result<Tag[]> raises <std::vision> {
  """
  Describe an image with booru tags using a local model.

  @param path - Path to a PNG, JPEG, WebP, or GIF image.
  @param model - A running local tagger, such as "wd14-tagger" or "florence-2".
  @param threshold - Minimum tag score from 0 to 1.
  @param limit - Maximum tags to return, from 1 to 500.
  """
  const file = try _visionFile(path)
  if (isFailure(file)) {
    return file
  }
  return interrupt std::vision("Tag this image with a local model?", {
    dir: file.value.dir,
    filename: file.value.filename,
    task: "tags",
    model: model,
  })
  return try _tagImage(file.value.image, model, threshold, limit)
}

@capabilities(fileRead, network)
export idempotent def captionImage(
  path: string,
  model: string,
  detail: string = "short",
): Result<string> raises <std::vision> {
  """
  Write a sentence about an image with a local model. Runs on this
  machine; nothing is uploaded.

  @param path - Path to a PNG, JPEG, WebP, or GIF image
  @param model - A local vision model that captions, such as "florence-2"
  @param detail - "short" for one sentence, "long" for a paragraph
  """
  const file = try _visionFile(path)
  if (isFailure(file)) {
    return file
  }
  return interrupt std::vision("Caption this image with a local model?", {
    dir: file.value.dir,
    filename: file.value.filename,
    task: "captions",
    model: model,
  })
  return try _captionImage(file.value.image, model, detail)
}

/** Compare vectors with `cosineSimilarity` from `std::embedding`. Use crops
from the same model and with similar framing. Compare against labeled
examples of both the objects you want and other objects; the module example
shows how to choose the closer group. */
@capabilities(fileRead, network)
export idempotent def embedImage(
  path: string,
  model: string,
  boxes: BoundingBox[] | null = null,
): Result<number[][]> raises <std::vision> {
  """
  Embed a whole image or selected regions for visual similarity comparisons.
  Compare vectors from the same model and with similar crop framing.

  @param path - Path to a PNG, JPEG, WebP, or GIF image.
  @param model - A running local embedding model, such as "dinov2-base".
  @param boxes - Up to 100 normalized regions, or null for the whole image.
  """
  const file = try _visionFile(path)
  if (isFailure(file)) {
    return file
  }
  return interrupt std::vision("Describe this image as numbers with a local model?", {
    dir: file.value.dir,
    filename: file.value.filename,
    task: "embeddings",
    model: model,
  })
  return try _embedImage(file.value.image, model, boxes)
}

/**
Use this before `embedImage` when you want candidate objects without choosing
category names. Regions arrive in descending score order. Florence-2 gives
every region a score of 1; OWLv2 defaults to a threshold of 0.1.
*/
@capabilities(fileRead, network)
export idempotent def findRegions(
  path: string,
  model: string,
  limit: number = 50,
  threshold: number | null = null,
): Result<Region[]> raises <std::vision> {
  """
  Find candidate object regions in an image using a local model, without labels.

  @param path - Path to a PNG, JPEG, WebP, or GIF image.
  @param model - A running local region model, such as "owlv2-base" or "florence-2".
  @param limit - Maximum regions to return, from 1 to 100.
  @param threshold - Minimum score from 0 to 1, or null for the model default.
  """
  const file = try _visionFile(path)
  if (isFailure(file)) {
    return file
  }
  return interrupt std::vision("Box the things in this image with a local model?", {
    dir: file.value.dir,
    filename: file.value.filename,
    task: "regions",
    model: model,
  })
  return try _findRegions(file.value.image, model, limit, threshold)
}