#!/usr/bin/env python3

import cv2

import numpy as np

from cvl.dataset import OnlineTrackingBenchmark
from cvl.trackers import NCCTracker
from cvl.grayscale_mosse import GrayscaleMosseTracker

dataset_path = "Mini-OTB"

SHOW_TRACKING = False
SAVE_IMAGES = True
SEQUENCE_IDX = 1
search_size = 1
learning_rate = 0.125
name="grayscale"
save_frame = 1

if __name__ == "__main__":

    dataset = OnlineTrackingBenchmark(dataset_path)

    a_seq = dataset[SEQUENCE_IDX]

    if SHOW_TRACKING:
        cv2.namedWindow("tracker")

    #tracker = NCCTracker()
    tracker = GrayscaleMosseTracker(save_img=SAVE_IMAGES,
                                    search_size=search_size,
                                    learning_rate=learning_rate,
                                    name=name,
                                    save_frame=save_frame)

    for frame_idx, frame in enumerate(a_seq):
        print(f"{frame_idx} / {len(a_seq)}")
        image_color = frame['image']

        if frame_idx == 0:
            bbox = frame['bounding_box']
            if bbox.width % 2 == 0:
                bbox.width += 1

            if bbox.height % 2 == 0:
                bbox.height += 1

            current_position = bbox
            tracker.start(image_color, bbox)
        else:
            tracker.detect(image_color)
            tracker.update(image_color)

        if SHOW_TRACKING:
            bbox = tracker.region
            pt0 = (bbox.xpos, bbox.ypos)
            pt1 = (bbox.xpos + bbox.width, bbox.ypos + bbox.height)
            image_color = cv2.cvtColor(image_color, cv2.COLOR_RGB2BGR)
            cv2.rectangle(image_color, pt0, pt1, color=(0, 255, 0), thickness=3)
            cv2.imshow("tracker", image_color)
            cv2.waitKey(0)
