from pathlib import Path

from ultralytics import YOLO
import cv2
from dodgeballtracker.helper import create_video_writer, iter_frames, time_tracker
from deep_sort_realtime.deepsort_tracker import DeepSort
import matplotlib
import numpy as np


Color = tuple[int, int, int]


class DetectedObjectBox:
    def __init__(self, min_threshold: float, color: Color):
        self.min_threshold = min_threshold
        self.color = color

    def __lt__(self, other):
        return self.min_threshold < other.min_threshold

    @classmethod
    def from_viridis(cls, min_threshold):
        cm = matplotlib.colormaps['viridis']
        color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3)))
        return cls(min_threshold, color)


def write_frame(writer, frame):
    cv2.imshow("Frame", frame)
    writer.write(frame)


def extract_box(data):
    return int(data[0]), int(data[1]), int(data[2]), int(data[3])


def plot_boxe_to_frame(data, frame, detected_box_objects):
    confidence = float(data[4])
    for detected_box_object in detected_box_objects:

        if detected_box_object.min_threshold < confidence:
            xmin, ymin, xmax, ymax = extract_box(data)
            cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
            return


def run_model_on_frame(model, frame):
    return model(frame)[0]


def crazy_experiment(frame):
    cv2.putText(frame,
                "Hello world",
                (50, 50),
                cv2.FONT_HERSHEY_SIMPLEX,
                2,
                (0, 255, 0), 8)
    cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5)
    pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32)
    pts = pts.reshape((-1, 1, 2))

    cv2.polylines(frame, pts, True, (0, 255, 255), 8)


def main(input_filename: Path, output_filename: Path, detected_box_objects):
    # define some constants
    detected_box_objects = sorted(detected_box_objects, reverse=True)

    video_cap = cv2.VideoCapture(str(input_filename))
    writer = create_video_writer(video_cap, str(output_filename))

    # load the pre-trained YOLOv8n model
    model = YOLO("yolov8n.pt")
    tracker = DeepSort(max_age=50)

    for frame in iter_frames(video_cap):
        with time_tracker():
            # run the YOLO model on the frame
            detections = run_model_on_frame(model, frame)
            for data in detections.boxes.data.tolist():
                plot_boxe_to_frame(data, frame, detected_box_objects)

            # show the frame to our screen
            write_frame(writer, frame)

            if cv2.waitKey(1) == ord("q"):
                break

    video_cap.release()
    writer.release()
    cv2.destroyAllWindows()