import cv2


class CanCounter:
    def __init__(self):
        self.all_pred = []
        self.count = 0
        self.max_val = 0
        self.length_list = []
        self.frame_count = 0
        self.send_frame = None

    def process_frame(self, frame, lis, info_center):
        pred = []
        for i in lis:
            conf = "{:.2f}".format(int(i['conf'] * 100))
            if int(i['conf'] * 100) > 60:
                label = 0
                try:
                    label = int(i['class'].item())
                except:
                    label = int(i['class'])
                det = i['points']
                if int(det[1]) > 208 and label == 0:
                    # print(i)
                    pred.append(i)
        if not pred and len(self.all_pred) > 0:
            self.all_pred.append(pred)
            self.length_list.append(0)
        for det in pred:
            conf = "{:.2f}".format(int(det['conf'] * 100))
            if int(det['conf'] * 100) > 60:
                label = 0
                try:
                    label = int(det['class'].item())
                except:
                    label = int(det['class'])
                det = det['points']

                if pred:
                    self.all_pred.append(pred)
                    self.length_list.append(int(len(pred)))
                    if label == 0:
                        cv2.rectangle(frame, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (0, 0, 255), 2)
                        print("done")
                    elif label == 1:
                        cv2.rectangle(frame, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (255, 255, 0), 2)
        length = len(self.all_pred)
        for i in range(len(self.all_pred)):
            if len(self.all_pred[i]) > self.max_val:
                self.send_frame = cv2.resize(frame, (64, 64))
            self.max_val = max(self.max_val, len(self.all_pred[i]))
            if length > 51:
                res = sum(self.length_list[-20:])
                if res == 0:
                    self.count = self.count + self.max_val
                    self.all_pred = []
                    if self.max_val in [9, 12, 15]:
                        box_count = info_center.get_box_count() + 1
                        info_center.send_payload(frame=self.send_frame,
                                                 message="Box: " + "BOX_" + str(box_count) + " | Cans: " + str(
                                                     self.max_val),
                                                 can_count=self.max_val, box_id="BOX_" + str(box_count),
                                                 eventtype="Completely Filled",
                                                 bg_color="#044b04", font_color="#24dc24")

                    else:
                        box_count = info_center.get_box_count() + 1
                        info_center.send_payload(frame=self.send_frame,
                                                 message="Box: " + "BOX_" + str(box_count) + " | Cans: " + str(
                                                     self.max_val),
                                                 can_count=self.max_val, box_id="BOX_" + str(box_count),
                                                 eventtype="Partially Filled",
                                                 bg_color="#472020", font_color="#ed2020", alert_sound="sound_1")

                    self.max_val = 0
                    break

        return frame
