# -*- coding: utf-8 -*-

from os import listdir
from numpy import ndarray
from cv2 import imread, findContours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE, imshow
from configparser import ConfigParser
import imageRecognition as IR
import getLicenseInformation as GLI
import imageProcess as IP


class MainProcess:
    """
    主进程类
    里面包括许多核心操作
    """

    def __init__(self):
        # 初始化 OCR
        self.ocr = IR.OCR()

        # 确定识别模式 1->默认 2->定位 如果选择模式2 后续会用到process_tool
        self.module = self.get_module()
        self.process_tool = None

        # 单张图片处理时 img_path 记录该图片路径
        self.img_path = ''
        self.img = None

        # 多张图片处理时 dir_path记录文件夹路径 下面两列表分别记录所有图片的文件名和路径
        self.dir_path = ''
        self.img_name_list = []
        self.img_path_list = []

        # 多张图片处理 以元组列表形式记录数据 例如[('1.jpg', '鲁A12345', '山东济南', '普通车辆'), (), ..]
        self.log_data = []

        # 车牌号定位器 初始化
        self.license_information = GLI.LicenseInformation()

    def set_img_path(self, path):
        self.img_path = path

    def set_dir_path(self, path):
        self.dir_path = path

    @staticmethod
    def get_module():
        config = ConfigParser()
        config.read('config.ini')
        module = int(dict(config.items('basic_settings'))['rec_module'])

        if module in (1, 2):
            return module
        return 1

    def get_dir_img_num(self):
        """获取文件夹内图片数量"""
        return len(self.img_path_list)

    def get_img_name_list(self):
        """获取待处理图片名称表"""
        return self.img_name_list

    def get_img_path_list(self):
        """获取待处理图片地址表"""
        return self.img_path_list

    @staticmethod
    def get_license_image(img):
        """处理图像 获取仅有车牌的图像"""
        if isinstance(img, str):
            img = imread(img)
        elif isinstance(img, ndarray):
            pass
        else:
            assert False, "识别出现异常"

        # 寻找车牌号 划分
        cut_image = IP.cutout_image(img)

        # 轮廓与层次
        contours, hierarchy = findContours(cut_image, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)

        # 切割出车牌号照片
        license_image = IP.get_license_image(contours, img)

        return license_image

    def get_dir_path_img(self, dir_path):
        """获取含有图片的文件夹  如果选择文件会出错"""
        file_list = []
        if dir_path != '':
            file_list = listdir(dir_path)

            self.img_name_list.clear()
            self.img_path_list.clear()

        for file in file_list:
            if file.endswith('.jpg') or file.endswith('.png'):
                self.img_name_list.append(file)
                self.img_path_list.append(dir_path + '/' + file)

    def process_single_image(self):
        try:
            # 获取路径并处理图像
            if self.module == 1:
                self.img = imread(self.img_path)
            elif self.module == 2:
                self.img = self.get_license_image(self.img_path)
                imshow('test', self.img)

            # 识别图像信息
            result = self.ocr.rec_image(self.img)

            # 分析结果
            if len(result):
                final_result = "找到了" + str(len(result)) + "个车牌号:\n"
                for i in range(len(result)):
                    final_result += ("车牌号" + str(i + 1) + ":" + result[i] + "\n")
                    final_result += ("归属地:" + self.license_information.get_location(result[i]) + "\n")
                    final_result += ("所属类型:" + self.license_information.get_type(result[i]) + "\n")
            else:
                final_result = "识别失败或本图片不包含车牌号,注意只支持国内车牌号如:鲁B12345"

            final_result += "\n"

        except:
            final_result = "出现问题:检查以下内容：\n1、图片路径不得有中文\n2、文件为jpg或png格式，正常无损坏\n3、内存分配和访问权限问题\n"

        return final_result

    def process_several_image(self, app=None, edit=None):
        """
        多张图片处理并记录结果
        如果app和edit不为None意味着是可视化界面在调用本方法 反之则是命令行界面调用
        """

        if self.module == 1:
            for img_index in range(len(self.img_name_list)):
                try:
                    result = "开始识别第%d张图片\n" % (img_index + 1)
                    img_path_i = self.img_path_list[img_index]
                    img_i = imread(img_path_i)
                    rec_result = self.ocr.rec_image(img_i)

                    if len(rec_result):
                        result += (self.img_name_list[img_index] + "这张图片中找到了" + str(len(rec_result)) + "个车牌号:\n")
                        for i in range(len(rec_result)):
                            result_t = rec_result[i]
                            location_t, type_t = self.license_information.get_location(
                                result_t), self.license_information.get_type(result_t)

                            result += ("车牌号" + str(i + 1) + ":" + result_t + "\n")
                            result += ("车牌号" + str(i + 1) + "归属地:" + location_t + "\n")
                            result += ("车牌号" + str(i + 1) + "所属类型:" + type_t + "\n")

                            self.log_data.append((self.img_name_list[img_index], result_t, location_t, type_t))
                    else:
                        result += "识别失败或本图片不包含车牌号,注意只支持国内车牌号如:鲁B12345"
                    result += '\n'

                    if app is not None and edit is not None:
                        edit.append(result)
                        if img_index % 5 == 0:
                            app.processEvents()
                    else:
                        print(result)
                except:
                    result = "出现问题,本次识别失败:检查以下内容\n1、文件夹路径不得有中文\n2、选择的是文件夹，程序会寻找文件夹内图片，支持jpg和png格式\n3、内存分配和访问权限问题\n"
                    if app is not None and edit is not None:
                        edit.append(result)
                        app.processEvents()
                    else:
                        print(result)

        elif self.module == 2:
            for img_index in range(len(self.img_name_list)):
                try:
                    result = "开始识别第%d张图片\n" % (img_index + 1)
                    img_path_i = self.img_path_list[img_index]
                    img_i = self.get_license_image(img_path_i)
                    rec_result = self.ocr.rec_image(img_i)

                    if len(rec_result):
                        result += (self.img_name_list[img_index] + "这张图片中找到了" + str(len(rec_result)) + "个车牌号:\n")
                        for i in range(len(rec_result)):
                            result_t = rec_result[i]
                            location_t, type_t = self.license_information.get_location(
                                result_t), self.license_information.get_type(result_t)

                            result += ("车牌号" + str(i + 1) + ":" + result_t + "\n")
                            result += ("车牌号" + str(i + 1) + "归属地:" + location_t + "\n")
                            result += ("车牌号" + str(i + 1) + "所属类型:" + type_t + "\n")

                            self.log_data.append((self.img_name_list[img_index], result_t, location_t, type_t))
                    else:
                        result += "识别失败或本图片不包含车牌号,注意只支持国内车牌号如:鲁B12345"
                    result += '\n'

                    if app is not None and edit is not None:
                        edit.append(result)
                        if img_index % 5 == 0:
                            app.processEvents()
                    else:
                        print(result)
                except:
                    result = "出现问题,本次识别失败:检查以下内容\n1、文件夹路径不得有中文\n2、选择的是文件夹，程序会寻找文件夹内图片，支持jpg和png格式\n3、内存分配和访问权限问题\n"
                    if app is not None and edit is not None:
                        edit.append(result)
                        app.processEvents()
                    else:
                        print(result)

        if app is not None and edit is not None:
            edit.append("识别完成!")
            app.processEvents()

        config = ConfigParser()
        config.read('config.ini')
        save_csv = True if (dict(config.items('basic_settings'))['save_csv']).lower() == 'true' else False
        if save_csv:
            self.save_data()

        return self.log_data

    def save_data(self):
        from csv import writer

        with open('data.csv', 'w', newline='') as f:
            writer = writer(f)
            writer.writerow(('源图片', '车牌号', '归属地', '车辆类型'))
            writer.writerows(tuple(self.log_data))