///|
pub struct MaskRun {
  value : Int
  length : Int
} derive(Debug, Eq)

///|
pub fn MaskRun::new(value~ : Int, length~ : Int) -> MaskRun {
  { value, length: max_i(0, length) }
}

///|
pub struct MaskRle {
  width : Int
  height : Int
  runs : Array[MaskRun]
} derive(Debug, Eq)

///|
pub fn Mask::to_rle(self : Mask) -> MaskRle {
  let runs : Array[MaskRun] = []
  if self.values.length() > 0 {
    let mut current = self.values[0]
    let mut length = 0
    for value in self.values {
      if value == current {
        length += 1
      } else {
        runs.push(MaskRun::new(value=current, length~))
        current = value
        length = 1
      }
    }
    runs.push(MaskRun::new(value=current, length~))
  }
  { width: self.width, height: self.height, runs }
}

///|
pub fn MaskRle::from_mask(mask : Mask) -> MaskRle {
  mask.to_rle()
}

///|
pub fn MaskRle::pixel_count(self : MaskRle, value : Int) -> Int {
  let mut total = 0
  for run in self.runs {
    if run.value == value {
      total += run.length
    }
  }
  total
}

///|
pub fn MaskRle::to_json(self : MaskRle) -> String {
  let sb = StringBuilder()
  sb.write_string("{\"size\":[\{self.height},\{self.width}],\"counts\":[")
  for i, run in self.runs {
    if i > 0 {
      sb.write_string(",")
    }
    sb.write_string("\{run.length}")
  }
  sb.write_string("]}")
  sb.to_string()
}

///|
pub struct CocoCategory {
  id : Int
  name : String
} derive(Debug, Eq)

///|
pub fn CocoCategory::new(id~ : Int, name~ : String) -> CocoCategory {
  { id, name }
}

///|
pub struct CocoImageRecord {
  id : Int
  file_name : String
  width : Int
  height : Int
} derive(Debug, Eq)

///|
pub fn CocoImageRecord::new(
  id~ : Int,
  file_name~ : String,
  width~ : Int,
  height~ : Int,
) -> CocoImageRecord {
  { id, file_name, width, height }
}

///|
pub fn CocoImageRecord::to_json(self : CocoImageRecord) -> String {
  "{\"id\":\{self.id},\"file_name\":\"\{escape_json(self.file_name)}\",\"width\":\{self.width},\"height\":\{self.height}}"
}

///|
pub struct CocoAnnotationRecord {
  id : Int
  image_id : Int
  category_id : Int
  bbox : Rect
  area : Int
  iscrowd : Int
} derive(Debug, Eq)

///|
pub fn CocoAnnotationRecord::from_box(
  id~ : Int,
  image_id~ : Int,
  box : Box2D,
) -> CocoAnnotationRecord {
  {
    id,
    image_id,
    category_id: box.label.id,
    bbox: box.rect,
    area: box.area(),
    iscrowd: if box.occluded {
      1
    } else {
      0
    },
  }
}

///|
pub fn CocoAnnotationRecord::to_json(self : CocoAnnotationRecord) -> String {
  let rect = self.bbox
  "{\"id\":\{self.id},\"image_id\":\{self.image_id},\"category_id\":\{self.category_id},\"bbox\":[\{rect.x},\{rect.y},\{rect.width},\{rect.height}],\"area\":\{self.area},\"iscrowd\":\{self.iscrowd}}"
}

///|
pub struct CocoDataset {
  description : String
  images : Array[CocoImageRecord]
  annotations : Array[CocoAnnotationRecord]
  categories : Array[CocoCategory]
} derive(Debug, Eq)

///|
pub fn CocoDataset::new(description~ : String) -> CocoDataset {
  { description, images: [], annotations: [], categories: [] }
}

///|
pub fn CocoDataset::add_image(
  self : CocoDataset,
  image : CocoImageRecord,
) -> Unit {
  self.images.push(image)
}

///|
pub fn CocoDataset::add_annotation(
  self : CocoDataset,
  annotation : CocoAnnotationRecord,
) -> Unit {
  self.annotations.push(annotation)
}

///|
pub fn CocoDataset::add_category(
  self : CocoDataset,
  category : CocoCategory,
) -> Unit {
  for existing in self.categories {
    if existing.id == category.id {
      return
    }
  }
  self.categories.push(category)
}

///|
fn[T] write_json_array(
  sb : StringBuilder,
  values : Array[T],
  write : (StringBuilder, T) -> Unit,
) -> Unit {
  sb.write_string("[")
  for i, value in values {
    if i > 0 {
      sb.write_string(",")
    }
    write(sb, value)
  }
  sb.write_string("]")
}

///|
pub fn CocoDataset::to_json(self : CocoDataset) -> String {
  let sb = StringBuilder()
  sb.write_string(
    "{\"info\":{\"description\":\"\{escape_json(self.description)}\"},",
  )
  sb.write_string("\"images\":")
  write_json_array(sb, self.images, fn(builder, image) {
    builder.write_string(image.to_json())
  })
  sb.write_string(",\"annotations\":")
  write_json_array(sb, self.annotations, fn(builder, annotation) {
    builder.write_string(annotation.to_json())
  })
  sb.write_string(",\"categories\":")
  write_json_array(sb, self.categories, fn(builder, category) {
    builder.write_string(
      "{\"id\":\{category.id},\"name\":\"\{escape_json(category.name)}\"}",
    )
  })
  sb.write_string("}")
  sb.to_string()
}

///|
pub fn SceneOutput::to_coco_json(self : SceneOutput) -> String {
  let dataset = CocoDataset::new(description="moonbit-synthetic-vision")
  dataset.add_image(
    CocoImageRecord::new(
      id=1,
      file_name=self.plan.name + ".pgm",
      width=self.image.width,
      height=self.image.height,
    ),
  )
  for box in self.boxes {
    dataset.add_category(
      CocoCategory::new(id=box.label.id, name=box.label.name),
    )
  }
  for i, box in self.boxes {
    dataset.add_annotation(
      CocoAnnotationRecord::from_box(id=i + 1, image_id=1, box),
    )
  }
  dataset.to_json()
}

///|
pub fn scene_mask_rle(scene : SceneOutput) -> String {
  scene.mask.to_rle().to_json()
}