from typing import NewType
from dataclasses import dataclass
import numpy as np
import pandera as pa
from pandera.dtypes import Float32, Float64

XCoord = NewType('XCoord', float)
YCoord = NewType('YCoord', float)
Score = NewType('Confidence', float)
Category = NewType('Category', float)


@dataclass
class Point:
    x: XCoord
    y: YCoord


VideoFrame = NewType('Frame', float)
VideoTimestamp = float


@dataclass
class DetectedObject:
    xmin: XCoord
    ymin: YCoord
    xmax: XCoord
    ymax: YCoord
    score: Score
    category: Category

    def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]:
        return (int(self.xmin), int(self.ymin)), (
            int(self.xmax), int(self.ymax))

    def as_nparray(self):
        return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score,
                         self.category])


def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema:
    return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols})


DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score",
                                      "category"])

TrackedDetectedObject_ = tensor_of(Float64,
                                   list(DetectedObject_.columns) + ["object_id", "index"])
# ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"])


Color = tuple[int, int, int]