import cv2
import base64
import traceback
from yolov5processor.infer import ExecuteInference
import time
import cv2
from edge_engine.common.logsetup import logger
from edge_engine.ai.model.modelwraper import ModelWrapper
import json
import os
from scripts.utils.infocenter import InfoCenter
from scripts.utils.cancounter import CanCounter

try:
    import urlparse
except ImportError:
    import urllib.parse as urlparse


class SRF_Cans(ModelWrapper):

    def __init__(self, config, pubs, device_id):
        super().__init__()
        """
        init function
        """
        self.config = config["config"]
        self.type = config['inputConf']['sourceType']
        with open('assests/cans.json', 'r') as f:
            self.dets = json.loads(f.read())
        self.rtp = pubs.rtp_write
        self.ic = InfoCenter(device_id=device_id)
        logger.info("Loading model to the device")
        self.can_counter_obj = CanCounter()
        self.base_model_path = 'assests/'
        print("[INFO] loading yolov5-Cans Detection Model")

        self.model_detector_pth = os.path.join(self.base_model_path, "cans_yolov5s.pt")
        self.yp = ExecuteInference(weight=self.model_detector_pth)
        self.pre_time = time.time()

    def _pre_process(self, x):
        """
        Do preprocessing here, if any
        :param x: payload
        :return: payload
        """
        return x

    def _post_process(self, x):
        """
        Apply post processing here, if any
        :param x: payload
        :return: payload
        """
        self.rtp.publish(x)  # video stream
        return x

    def _predict(self, x):
        try:
            time.sleep(0.02)
            frame = x['frame']
            frame_id = x['frameId']
            frame = cv2.resize(frame, (640, 400))
            lis = self._type(self.type, frame, frame_id)
            frame = self.can_counter_obj.process_frame(frame, lis=lis, info_center=self.ic)
            curr_time = time.time()
            logger.info("Time for 1 frame: {}".format(curr_time - self.pre_time))
            self.pre_time = curr_time
            x['frame'] = cv2.resize(frame, (self.config.get('FRAME_WIDTH'), self.config.get('FRAME_HEIGHT')))
            return x

        except Exception as e:
            logger.error(f"Error: {e}", exc_info=True)
            x['frame'] = cv2.resize(x['frame'], (self.config.get('FRAME_WIDTH'), self.config.get('FRAME_HEIGHT')))
            return x

    def _type(self, video_type, frame, frame_id):
        if video_type not in ['videofile']:
            return self.yp.predict(frame)
        print("Running without")
        return self.dets[int(frame_id)][str(frame_id)]['detections']
