from enum import Enum

### Configuration options
WEIGHTS = 'models/v5m_896_300best.pt'
# will need to configure these based on GPU hardware
BATCH_SIZE = 32

CONF_THRES = 0.05 # detection
NMS_IOU  = 0.2 # NMS IOU
MAX_AGE = 14 # time until missing fish get's new id
MIN_HITS = 16 # minimum number of frames with a specific fish for it to count
MIN_LENGTH = 0.3 # minimum fish length, in meters
IOU_THRES = 0.01 # IOU threshold for tracking
MIN_TRAVEL = -1 # Minimum distance a track has to travel

class TrackerType(Enum):
    NONE = 0
    CONF_BOOST = 1
    BYTETRACK = 2

class InferenceConfig:
    def __init__(self, 
                weights=WEIGHTS, conf_thresh=CONF_THRES, nms_iou=NMS_IOU, 
                min_hits=MIN_HITS, max_age=MAX_AGE, min_length=MIN_LENGTH, min_travel=MIN_TRAVEL):
        self.weights = weights
        self.conf_thresh = conf_thresh
        self.nms_iou = nms_iou
        self.min_hits = min_hits
        self.max_age = max_age
        self.min_length = min_length
        self.min_travel = min_travel

        self.associative_tracker = TrackerType.NONE
        self.boost_power = 1
        self.boost_decay = 1
        self.byte_low_conf = 1
        self.byte_high_conf = 1

    def enable_conf_boost(self, power, decay):
        self.associative_tracker = TrackerType.CONF_BOOST
        self.boost_power = power
        self.boost_decay = decay

    def enable_byte_track(self, low, high):
        self.associative_tracker = TrackerType.BYTETRACK
        self.byte_low_conf = low
        self.byte_high_conf = high

    def to_dict(self):
        dict = {
            'weights': self.weights,
            'nms_iou': self.nms_iou,
            'min_hits': self.min_hits,
            'max_age': self.max_age,
            'min_length': self.min_length,
            'min_travel': self.min_travel,
        }

        # Add tracker specific parameters
        if (self.associative_tracker == TrackerType.BYTETRACK):
            dict['tracker'] = "ByteTrack"
            dict['byte_low_conf'] = self.byte_low_conf
            dict['byte_high_conf'] = self.byte_high_conf
        elif (self.associative_tracker == TrackerType.CONF_BOOST):
            dict['tracker'] = "Confidence Boost"
            dict['conf_thresh'] = self.conf_thresh
            dict['boost_power'] = self.boost_power
            dict['boost_decay'] = self.boost_decay
        elif (self.associative_tracker == TrackerType.NONE):
            dict['tracker'] = "None"
            dict['conf_thresh'] = self.conf_thresh

        return dict