import cv2
import base64
import numpy as np
from scipy.spatial import distance
from expiringdict import ExpiringDict
import time
from edge_engine.common.logsetup import logger
from scripts.utils.infocenter import MongoLogger
from yolov5processor.infer import ExecuteInference
from scripts.utils.edge_utils import get_extra_fields
from edge_engine.ai.model.modelwraper import ModelWrapper
from scripts.utils.centroidtracker import CentroidTracker
from scripts.common.constants import JanusDeploymentConstants
from scripts.utils.image_utils import draw_circles_on_frame, resize_to_64_64
from pymongo import MongoClient
from scripts.common.config import MONGO_URI
from uuid import uuid4
import cv2
import base64
import datetime
import numpy as np
import imutils
from collections import deque
from expiringdict import ExpiringDict
from scipy.optimize import linear_sum_assignment as linear_assignment

from edge_engine.common.logsetup import logger
from edge_engine.ai.model.modelwraper import ModelWrapper

from scripts.utils.tracker import Tracker
from scripts.utils.helpers import box_iou2
from scripts.utils.edge_utils import Utilities
from scripts.utils.infocenter import MongoLogger
from scripts.utils.model_tracker import ModelCountTracker
from scripts.common.constants import JanusDeploymentConstants

from yolov5processor.infer import ExecuteInference
from scripts.utils.relay_util import RelayHandler
from paddleocr import PaddleOCR


class CementBagCounter(ModelWrapper):
    def __init__(self, config, model_config, pubs, device_id):
        super().__init__()
        """
        init function
        """
        self.config = config["config"]
        self.device_id = device_id
        self.rtp = pubs.rtp_write
        self.mongo_logger = MongoLogger()
        self.frame_skip = self.config.get("frame_skip", False)
        # model = "data/acc_v13.pt"
        # self.yp = ExecuteInference(
        #     weight=model,
        #     gpu=model_config.get("gpu", False),
        #     agnostic_nms=model_config.get("agnostic_nms", True),
        #     iou=model_config.get("iou", 0.2),
        #     confidence=model_config.get("confidence", 0.3),
        #     img_size=640,
        # )
        # self.print_eu_dist = model_config.get("print_eu_dist", 200)
        # self.ct1 = CentroidTracker(maxDisappeared=5)
        # self.ct2 = CentroidTracker(maxDisappeared=5)

        self.count = 0
        self.cement_bag = 0
        self.count_suraksha = 0
        self.count_whitecem = 0
        self.count_gold = 0
        self.tracker_list = []
        self.max_age = 3
        self.min_hits = 0
        self.track_id_list = deque([str(i) for i in range(1, 50)])
        self.prev_annotation = []
        self.mrp_counter = 0
        self.count_nfr = 0
        self.count_suraksha_power = 0
        self.count_concrete_plus = 0
        self.count_ambuja_plus = 0
        self.text = " "
        self.mrp_text = " "
        self.recognition_output_list = []
        self.recognition_output=""
        # self.prev_class_name = None
        self.plant_name = ["R"]
        self.year_front_axle_bar_and_engine = ["B","C","D","E","F","G","H","I","J","K","L","M","N","O","P","Q","R","S","T","U","V","W","X","Y","Z","1","2","3","4"]
        self.fixed = ["M","B","N"]
        self.internal_engine = ["2", "K", "G", "A"]
        self.internal =["T","F","A","E","X","A"]
        self.internal_front_axle_bar_2  = ["A","B","A","C","A","A"]
        self.internal_front_axle_bar = ["T","F","A","E","X","A"]
        # self.fixed_and_internal_front_axle_bar_2 = ["M", "B", "N", "A", "B", "A", "C", "A", "A"]
        self.fixed_and_internal_engine = ["2", "K", "C", "A"]
        self.month_engine = ["K","L","M","A","B","C","D","E","F","G","H","J"]
        self.month_front_axle_bar = ["A","B","C","D","E","F","G","H","I","J","K","L","M"]
        self.text_json = {"data":{"text":" ", "fixed characters":" ", "internal characters": " ", "year": " ", "month": " ", "plant": " ", "serial number": " "}, "status": False, "message": "Image not clear"}
        self.digit_list = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0"]
        self.initial_object_position = None

        self.uncounted_objects = ExpiringDict(
            max_len=model_config.get("uncounted_obj_length", 50),
            max_age_seconds=model_config.get("uncounted_obj_age", 60),
        )
        self.janus_metadata = ExpiringDict(max_age_seconds=120, max_len=1)
        self.mongo_alarm_coll = MongoClient(MONGO_URI)["ilens_events"][
            "triggered_alarms"
        ]
        self.camera_details = self.mongo_logger.get_camera_details(self.device_id)
        self.ocr=self.paddle_ocr_load_model()
        # self.black_to_white_dict = get_extra_fields(self.device_id).get(
        #     JanusDeploymentConstants.BLACK_WHITE_RATIO
        # )


    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 send_payload(
        self,
        frame,
        bag_type="",
        label="CementBagDetected",
        bg_color="#474520",
        font_color="#FFFF00",
        alert_sound=None,
        message="Cement Bag Detected!",
        mrp="FAIL",
        mrp_frmae="",
        mrp_roi="",
        mrp_check_result="",
    ):
        """
        Insert event to Mongo
        :param message:
        :param frame:
        :param label:
        :param bg_color:
        :param font_color:
        :param alert_sound:
        :return: None
        """

        payload = {
            "deviceId": self.device_id,
            "message": message,
            "frame": "data:image/jpeg;base64,"
            + base64.b64encode(cv2.imencode(".jpg", frame)[1].tostring()).decode(
                "utf-8"
            ),
            "activity": label,
            "bg_color": bg_color,
            "font_color": font_color,
            "alert_sound": alert_sound,
            "bag_type": bag_type,
            "mrp_frmae": mrp_frmae,
            "mrp_roi": mrp_roi,
            "mrp_digit_check": mrp_check_result,
            "mrp": mrp
        }

        self.mongo_logger.insert_attendance_event_to_mongo(payload)

    def paddle_ocr_load_model(self):
        ocr = PaddleOCR(
            lang="en",
            # det_db_thresh=0.1,
            # det_db_box_thresh=0.1,
            # use_mp=True,
            # total_process_num=process_count,
            use_angle_cls=True,
            cls_model_dir="paddleocr/model/ch_ppocr_mobile_v2.0_cls_infer",
            rec_model_dir="paddleocr/model/ch_PP-OCRv3_rec_infer",
            det_model_dir="paddleocr/model/en_PP-OCRv3_det_infer")
        return ocr
    def paddle_ocr_predict(self,ocr,img_path):

        result = ocr.ocr(img_path, cls=False, det=True, rec=True)[0]
        txts = [line[1][0] for line in result]
        return txts[0]

    def check_character(self,character):
        if len(character) == 17:
            return "front_axle_bar"
        elif len(character) == 13:
            return "engine"

        else:
            return False
    def fixed_char_check(self, text_list, part_name):
        if(part_name == "front_axle_bar"):
            sub_list = text_list[0:3]
            # print(sub_list)
            matching_chars = [i for i, j in zip(sub_list, self.fixed) if i == j]
            # matching_chars_2 = [i for i, j in zip(sub_list, self.fixed_and_internal_front_axle_bar_2) if i == j]
            # print(matching_chars)
            if(len(matching_chars) > 1):
                return True
            return False

        else:

            return True
    def internal_char_check(self, text_list, part_name):
        if(part_name == "front_axle_bar"):
            sub_list = text_list[3:9]
            # print(sub_list)
            matching_chars = [i for i, j in zip(sub_list, self.internal_front_axle_bar) if i == j]
            matching_chars_2 = [i for i, j in zip(sub_list, self.internal_front_axle_bar_2) if i == j]
            # print(matching_chars)
            if(len(matching_chars) > 3):
                return True, "front axle_bar"
            if (len(matching_chars_2) > 3):
                return True, "front axle bar 2"
            return False, None
        else:
            sub_list = text_list[3:7]
            # print(sub_list)
            matching_chars = [i for i, j in zip(sub_list, self.internal_engine) if i == j]
            # print(matching_chars)
            if (len(matching_chars) > 0):
                # print("1234567890!@#$%^&*()")
                return True, None
            return False, None

    def plant_ver(self, text_list, part_name):
        if (part_name == "front_axle_bar"):

            plant_char = text_list[10]
            if (plant_char != "R"):
                return False
            return True

        else:
            plant_char = text_list[0]
            if (plant_char != "R"):
                return False
            return True
    def year_ver(self, text_list, part_name):
        if (part_name == "front_axle_bar"):

            year_char = text_list[9]
            if (year_char not in self.year_front_axle_bar_and_engine):
                return False
            return True

        else:
            year_char = text_list[1]
            if (year_char not in self.year_front_axle_bar_and_engine):
                return False
            return True
    def month_ver(self, text_list, part_name):
        if (part_name == "front_axle_bar"):

            month_char = text_list[11]
            if (month_char not in self.month_front_axle_bar):
                return False
            return True

        else:
            month_char = text_list[2]
            if (month_char not in self.month_engine):
                return False
            return True
    def digit_correction(self, text_list, part_name):
        corrected_serial_number = []
        if (part_name == "front_axle_bar"):

            serial_number = text_list[12:]
            for ind, elem in enumerate(serial_number):
                if(elem == "O" or elem == "o"):
                    corrected_serial_number.append("0")
                elif(elem == "?" or elem == "T"):
                    corrected_serial_number.append("7")
                elif(elem == "U"):
                    corrected_serial_number.append("4")
                elif(elem == "I" or elem == "i"):
                    corrected_serial_number.append("1")
                else:
                    corrected_serial_number.append(elem)
        else:

            serial_number = text_list[7:]
            for ind, elem in enumerate(serial_number):
                if(elem == "O" or elem == "o"):
                    corrected_serial_number.append("0")
                elif(elem == "?" or elem == "T"):
                    corrected_serial_number.append("7")
                elif(elem == "U"):
                    corrected_serial_number.append("4")
                elif(elem == "I" or elem == "i"):
                    corrected_serial_number.append("1")
                else:
                    corrected_serial_number.append(elem)

        return corrected_serial_number

    def defining_actual_text(self, text_list, part_name, fixed_verification, internal_verification, chassis_part, plant_verification, year_verification, month_verification, corrected_digit):
        actual_text = []
        if(part_name == "front_axle_bar"):
            if(fixed_verification):
                actual_text.extend(self.fixed)
            else:
                return "Image not clear", False
            if(internal_verification):
                if(chassis_part == "front axle_bar"):
                    actual_text.extend(self.internal_front_axle_bar)
                else:
                    actual_text.extend(self.internal_front_axle_bar_2)

            if(year_verification):
                actual_text.append(text_list[9])
            else:
                return "Image not clear", False
            if(month_verification):
                actual_text.append(text_list[11])
            else:
                return "Image not clear", False
            if(plant_verification):
                actual_text.append("R")
            else:
                return "Image not clear", False
            for elem in corrected_digit:
                print("elements are----")
                print(elem)
                if elem not in self.digit_list:
                    return "Image not clear", False

            actual_text.extend(corrected_digit)
        else:
            if (plant_verification):
                actual_text.append("R")
            else:
                return "Image not clear", False
            if (year_verification):
                actual_text.append(text_list[1])
            else:
                return "Image not clear", False

            if (month_verification):
                actual_text.append(text_list[2])
            else:
                return "Image not clear", False

            if (internal_verification):

                actual_text.extend(self.internal_engine)


            if (fixed_verification):
                pass
            else:
                return "Image not clear", False



            for elem in corrected_digit:
                print("elements are----")
                print(elem)
                if elem not in self.digit_list:
                    return "Image not clear", False


            actual_text.extend(corrected_digit)
        return actual_text, True


    def _predict(self, obj):

        try:
            frame = obj["frame"]
            text = self.paddle_ocr_predict(self.ocr,frame)
            print("detected text is ---------------", text)
            part_name =self.check_character(text)
            if not part_name:
                text="Image not Clear"
                self.text_json = {
                    "data": {"text": " ", "fixed characters": " ", "internal characters": " ", "year": " ",
                             "month": " ", "plant": " ", "serial number": " "}, "status": False,
                    "message": "Image not clear"}
                pass
            else:
                text_list = []
                for letter in text:
                    text_list.append(letter)
                # print(text_list)
                fixed_verification = self.fixed_char_check(text_list, part_name)
                internal_verification, chassis_part = self.internal_char_check(text_list, part_name)
                # if char_verification1:
                #     self.recognition_output_list.extend(self.fixed_and_internal_trackter)
                # else:
                #     text = "Image not Clear"

                # if(plant_verification):
                plant_verification = self.plant_ver(text_list, part_name)
                year_verification = self.year_ver(text_list, part_name)
                month_verification = self.month_ver(text_list, part_name)
                corrected_digit = self.digit_correction(text_list, part_name)
                redefined_text, status = self.defining_actual_text(text_list, part_name, fixed_verification, internal_verification, chassis_part, plant_verification, year_verification, month_verification, corrected_digit)
                if(status):
                    if(part_name == "front_axle_bar"):
                        year = " "
                        month = " "

                        fixed_char = redefined_text[0:3]
                        internal_char = redefined_text[3:9]
                        year_char = redefined_text[9]
                        plant_char = redefined_text[11]
                        month_char = redefined_text[10]
                        print("MONTH------------------")
                        print(month_char)
                        serial_number_char = redefined_text[12:]
                        if(year_char == "N"):
                            year = "2022"
                        if(month_char == "H"):
                            month = "Augest"
                        elif(month_char == "C"):
                            month = "March"
                        plant = "M$M"
                        print("****************************************************************")
                        internal = ''.join(str(x) for x in internal_char)
                        fixed = ''.join(str(x) for x in fixed_char)
                        serial_number = ''.join(str(x) for x in serial_number_char)
                        self.recognition_output = ''.join(str(x) for x in redefined_text)
                        self.text_json["data"]["text"] = self.recognition_output
                        self.text_json["data"]["fixed characters"] = fixed
                        self.text_json["data"]["internal characters"] = internal
                        self.text_json["data"]["year"] = year
                        self.text_json["data"]["month"] = month
                        self.text_json["data"]["plant"] = plant
                        self.text_json["data"]["serial number"] = serial_number
                        self.text_json["status"] = True
                        self.text_json["message"] = " "
                    else:
                        print("engine part---------------------")
                        year = " "
                        month = " "
                        # fixed_char = redefined_text[0:3]
                        internal_char = redefined_text[3:7]
                        year_char = redefined_text[1]
                        plant_char = redefined_text[0]
                        month_char = redefined_text[2]
                        serial_number_char = redefined_text[7:]
                        if (year_char == "N"):
                            year = "2022"
                        if (month_char == "F"):
                            month = "September"
                        plant = "M$M"
                        internal = ''.join(str(x) for x in internal_char)
                        fixed = " "
                        serial_number = ''.join(str(x) for x in serial_number_char)

                    self.recognition_output = ''.join(str(x) for x in redefined_text)
                    self.text_json["data"]["text"] = self.recognition_output
                    self.text_json["data"]["fixed characters"] = fixed
                    self.text_json["data"]["internal characters"] = internal
                    self.text_json["data"]["year"] = year
                    self.text_json["data"]["month"] = month
                    self.text_json["data"]["plant"] = plant
                    self.text_json["data"]["serial number"] = serial_number
                    self.text_json["status"] = True
                    self.text_json["message"] = " "

                else:

                    self.text_json["data"]["fixed characters"] = " "
                    self.text_json["data"]["internal characters"] = " "
                    self.text_json["data"]["year"] = " "
                    self.text_json["data"]["month"] = " "
                    self.text_json["data"]["plant"] = " "
                    self.text_json["data"]["serial number"] = " "

                    # self.recognition_output = "Image not clear"
                    self.text_json["data"]["text"] = " "
                    self.text_json["status"] = False
                    self.text_json["message"] = "Image is not clear"

                print(self.text_json)


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