import abc
from pathlib import Path

import pandas as pd
import pandera as pa

from dodgeballtracker.typing import (VideoFrame, Color, TrackedDetectedObject)
import matplotlib
import cv2


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 create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter:

    # grab the width, height, and fps of the frames in the video stream.
    frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(video_cap.get(cv2.CAP_PROP_FPS))

    # initialize the FourCC and a video writer object
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(str(output_filename.absolute()),
                             fourcc,
                             fps,
                             (frame_width, frame_height))

    return writer


class VideoWriter(abc.ABC):

    @abc.abstractmethod
    def write_frame(self, frame: VideoFrame) -> None:
        pass

    @abc.abstractmethod
    def plot_boxes_to_frame(self, data: pd.DataFrame, frame: VideoFrame) -> None:
        pass

    @abc.abstractmethod
    def close(self) -> None:
        pass

    @abc.abstractmethod
    def wait(self) -> bool:
        pass


class NoopVideoWriter(VideoWriter):

    def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None:
        pass

    def write_frame(self, frame: VideoFrame) -> None:
        pass

    def close(self) -> None:
        pass

    def wait(self) -> bool:
        return False


class MP4VVideoWriter(VideoWriter):

    def __init__(self, video_writer: cv2.VideoWriter, thresholds: list[DetectedObjectBox]) -> None:
        self._video_writer = video_writer
        self._thresholds = thresholds

    def write_frame(self, frame: VideoFrame) -> None:
        self._video_writer.write(frame)

    def write_box_id(self, frame: VideoFrame, box: pd.Series):
        bottom_left_corner_of_text = (int(box['xmin']), int(box['ymax']))
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        font_color = (255, 255, 255)
        thickness = 1
        line_type = 2
        id_ = str(int(box['object_id']))
        cv2.putText(frame,
                    id_,
                    bottom_left_corner_of_text,
                    font,
                    font_scale,
                    font_color,
                    thickness,
                    line_type)

    def plot_box_to_frame(self, data: pd.Series, frame: VideoFrame) -> None:
        confidence = data['score']
        for threshold in self._thresholds:
            if threshold.min_threshold < confidence:
                start_point =  data[['xmin', 'ymin']].array.astype('int32')
                end_point = data[['xmax', 'ymax']].array.astype('int32')
                cv2.rectangle(frame, start_point, end_point, threshold.color, 2)
                self.write_box_id(frame, data)
                return

    @pa.decorators.check_io(boxes=TrackedDetectedObject)
    def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None:
        boxes_that_passe_threshold = boxes[boxes['score'] > 0.5]
        for i, box in boxes_that_passe_threshold.iterrows():
            self.plot_box_to_frame(box, frame)


    def close(self) -> None:
        self._video_writer.release()
        cv2.destroyAllWindows()

    def wait(self) -> bool:
        return False

    @classmethod
    def from_frame(cls, input_: Path, output: Path, thresholds: list[float]) -> VideoWriter:
        object_drawers = [DetectedObjectBox.from_viridis(t) for t in thresholds]
        video_cap = cv2.VideoCapture(str(input_))
        try:
            return cls(create_video_writer(video_cap, output), object_drawers)
        finally:
            video_cap.release()


class MP4VVideoWriterWithFrame(MP4VVideoWriter):
    def write_frame(self, frame: VideoFrame) -> None:
        super().write_frame(frame)
        cv2.imshow("Frame", frame)

    def wait(self) -> bool:
        return cv2.waitKey(1) == ord("q")