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

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])


DetectedObject_ = pa.DataFrameSchema({
    "xmin": pa.Column(Float32),
    "ymin": pa.Column(Float32),
    "xmax": pa.Column(Float32),
    "ymin": pa.Column(Float32),
    "score": pa.Column(Float32),
    "category": pa.Column(Float32),
})

TrackedDetectedObject_ = pa.DataFrameSchema({
    "xmin": pa.Column(Float32),
    "ymin": pa.Column(Float32),
    "xmax": pa.Column(Float32),
    "ymin": pa.Column(Float32),
    "object_id": pa.Column(Float32),
    "score": pa.Column(Float32),
    "category": pa.Column(Float32),
    "index": pa.Column(Float32),
})


Color = tuple[int, int, int]