///|
pub enum TrackState {
  Tentative
  Confirmed
  Lost
  Finished
} derive(Debug, Eq)

///|
pub struct TrackObservation {
  frame : Int
  detection : Box2D
  score_per_1000 : Int
} derive(Debug, Eq)

///|
pub fn TrackObservation::new(
  frame~ : Int,
  detection~ : Box2D,
  score_per_1000? : Int = 1000,
) -> TrackObservation {
  { frame, detection, score_per_1000: clamp(score_per_1000, low=0, high=1000) }
}

///|
pub struct Track {
  id : Int
  label : Label
  mut state : TrackState
  observations : Array[TrackObservation]
  mut missed_frames : Int
} derive(Debug, Eq)

///|
pub fn Track::new(id~ : Int, observation : TrackObservation) -> Track {
  {
    id,
    label: observation.detection.label,
    state: Tentative,
    observations: [observation],
    missed_frames: 0,
  }
}

///|
pub fn Track::last(self : Track) -> TrackObservation? {
  if self.observations.length() == 0 {
    None
  } else {
    Some(self.observations[self.observations.length() - 1])
  }
}

///|
pub fn Track::age(self : Track) -> Int {
  if self.observations.length() == 0 {
    0
  } else {
    self.observations[self.observations.length() - 1].frame -
    self.observations[0].frame +
    1
  }
}

///|
pub fn Track::center(self : Track) -> Point {
  match self.last() {
    None => Point::new(x=0, y=0)
    Some(observation) => observation.detection.center()
  }
}

///|
pub fn Track::velocity(self : Track) -> Point {
  if self.observations.length() < 2 {
    Point::new(x=0, y=0)
  } else {
    let last = self.observations[self.observations.length() - 1].detection.center()
    let previous = self.observations[self.observations.length() - 2].detection.center()
    Point::new(x=last.x - previous.x, y=last.y - previous.y)
  }
}

///|
pub fn Track::add(
  self : Track,
  observation : TrackObservation,
  confirmation_frames : Int,
) -> Unit {
  self.observations.push(observation)
  self.missed_frames = 0
  if self.observations.length() >= max_i(1, confirmation_frames) {
    self.state = Confirmed
  }
}

///|
pub fn Track::mark_missed(self : Track, max_missed : Int) -> Unit {
  self.missed_frames += 1
  if self.missed_frames > max_i(0, max_missed) {
    self.state = Finished
  } else {
    self.state = Lost
  }
}

///|
pub struct TrackerConfig {
  iou_threshold_per_1000 : Int
  max_missed_frames : Int
  confirmation_frames : Int
} derive(Debug, Eq)

///|
pub fn TrackerConfig::new(
  iou_threshold_per_1000? : Int = 300,
  max_missed_frames? : Int = 2,
  confirmation_frames? : Int = 2,
) -> TrackerConfig {
  {
    iou_threshold_per_1000: clamp(iou_threshold_per_1000, low=0, high=1000),
    max_missed_frames: max_i(0, max_missed_frames),
    confirmation_frames: max_i(1, confirmation_frames),
  }
}

///|
pub struct Tracker {
  config : TrackerConfig
  mut next_id : Int
  tracks : Array[Track]
} derive(Debug, Eq)

///|
pub fn Tracker::new(config? : TrackerConfig = TrackerConfig::new()) -> Tracker {
  { config, next_id: 1, tracks: [] }
}

///|
pub fn Tracker::active(self : Tracker) -> Array[Track] {
  let out : Array[Track] = []
  for track in self.tracks {
    if track.state != Finished {
      out.push(track)
    }
  }
  out
}

///|
pub fn Tracker::update(
  self : Tracker,
  frame : Int,
  detections : Array[Box2D],
) -> Array[Track] {
  let used = Array::make(detections.length(), false)
  for track in self.tracks {
    if track.state == Finished {
      continue
    }
    let mut best = -1
    let mut best_iou = self.config.iou_threshold_per_1000 - 1
    match track.last() {
      None => ()
      Some(previous) =>
        for i, detection in detections {
          if !used[i] && detection.label.id == track.label.id {
            let score = previous.detection.iou(detection)
            if score > best_iou {
              best = i
              best_iou = score
            }
          }
        }
    }
    if best >= 0 {
      used[best] = true
      track.add(
        TrackObservation::new(frame~, detection=detections[best]),
        self.config.confirmation_frames,
      )
    } else {
      track.mark_missed(self.config.max_missed_frames)
    }
  }
  for i, detection in detections {
    if !used[i] {
      let track = Track::new(
        id=self.next_id,
        TrackObservation::new(frame~, detection~),
      )
      self.next_id += 1
      self.tracks.push(track)
    }
  }
  self.active()
}

///|
pub fn Tracker::finished(self : Tracker) -> Array[Track] {
  let out : Array[Track] = []
  for track in self.tracks {
    if track.state == Finished {
      out.push(track)
    }
  }
  out
}

///|
pub struct TrackSummary {
  id : Int
  label_id : Int
  age : Int
  state : TrackState
  center : Point
  velocity : Point
} derive(Debug, Eq)

///|
pub fn TrackSummary::from_track(track : Track) -> TrackSummary {
  {
    id: track.id,
    label_id: track.label.id,
    age: track.age(),
    state: track.state,
    center: track.center(),
    velocity: track.velocity(),
  }
}

///|
pub fn tracker_summaries(tracker : Tracker) -> Array[TrackSummary] {
  let out : Array[TrackSummary] = []
  for track in tracker.active() {
    out.push(TrackSummary::from_track(track))
  }
  out
}

///|
pub fn tracks_to_boxes(tracks : Array[Track]) -> Array[Box2D] {
  let out : Array[Box2D] = []
  for track in tracks {
    match track.last() {
      None => ()
      Some(observation) => out.push(observation.detection)
    }
  }
  out
}

///|
pub fn predict_box(track : Track, frames : Int) -> Box2D? {
  match track.last() {
    None => None
    Some(observation) => {
      let velocity = track.velocity()
      let offset = Point::new(
        x=velocity.x * max_i(0, frames),
        y=velocity.y * max_i(0, frames),
      )
      Some(
        Box2D::new(
          label=observation.detection.label,
          rect=observation.detection.rect.translate(offset.x, offset.y),
          occluded=observation.detection.occluded,
        ),
      )
    }
  }
}

///|
pub fn tracker_to_json(tracker : Tracker) -> String {
  let summaries = tracker_summaries(tracker)
  let sb = StringBuilder()
  sb.write_string("[ ")
  for i, summary in summaries {
    if i > 0 {
      sb.write_string(",")
    }
    sb.write_string(
      "{\"id\":\{summary.id},\"label\":\{summary.label_id},\"age\":\{summary.age},\"state\":\"\{track_state_name(summary.state)}\",\"x\":\{summary.center.x},\"y\":\{summary.center.y}}",
    )
  }
  sb.write_string(" ]")
  sb.to_string()
}

///|
pub fn track_state_name(state : TrackState) -> String {
  match state {
    Tentative => "tentative"
    Confirmed => "confirmed"
    Lost => "lost"
    Finished => "finished"
  }
}

///|
pub fn validate_tracker_config(config : TrackerConfig) -> ValidationReport {
  let report = ValidationReport::new()
  if config.iou_threshold_per_1000 < 0 || config.iou_threshold_per_1000 > 1000 {
    report.error("tracker IoU threshold must be between 0 and 1000")
  }
  if config.confirmation_frames < 1 {
    report.error("tracker confirmation frames must be positive")
  }
  report
}