#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File        :   dot_detection.py
@Contact     :   zpyovo@hotmail.com
@License     :   (C)Copyright 2018-2019, Lab501-TransferLearning-SCUT
@Description :

@Modify Time      @Author    @Version    @Desciption
------------      -------    --------    -----------
2022/4/21 17:53   Pengyu Zhang      1.0         None
'''

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cv2
from common import automatedMSRCR, save_figure
from component.utile import return_img_stream
import numpy as np
import math
'''
采用realesrgan模型实现图像增强后进行网线角度检测
'''
def dots_angle_measure_dl_sr(dots_region,realesrgan,save_dots=False,res_code=None):

    original_img_stream = None
    process_dots_img_stream = None
    detection_dots_img_stream = None
    lines_arctan = None
    if save_dots:
        result_path = save_figure(dots_region, 'original', res_code)
        original_img_stream = return_img_stream(result_path)
    # 超分计算处理
    process_dots = realesrgan(dots_region)


    if save_dots:
        b, g, r = cv2.split(process_dots)
        process_dots_ = cv2.merge([r, g, b])
        result_path = save_figure(process_dots_, 'process_dots', res_code)
        process_dots_img_stream = return_img_stream(result_path)

    # 斑点检测
    found, corners = cv2.findCirclesGrid(process_dots, (3, 3), cv2.CALIB_CB_SYMMETRIC_GRID)

    if corners is None or corners.any() != None and corners.shape[0] < 2:
        # print("---------------------------------------------------------------")
        # print("corners status,", found)
        # print("---------------------------------------------------------------")
        return lines_arctan, [original_img_stream, process_dots_img_stream]

    elif corners.shape[0] >= 2:
        A1 = corners[0]
        A2 = corners[1]
        if (A1 == A2).all():
            # 斑点检测
            found, corners = cv2.findCirclesGrid(process_dots, (4, 4), cv2.CALIB_CB_SYMMETRIC_GRID)
            A1 = corners[0]
            A2 = corners[1]
            if (A1 == A2).all():
                return lines_arctan, [original_img_stream, process_dots_img_stream]
        B1 = corners[1]
        B2 = np.array([[0, corners[1][:, 1]]],dtype=object)

        kLine1 = (A2[:, 1] - A1[:, 1]) / (A2[:, 0] - A1[:, 0])
        kLine2 = (B2[:, 1] - B1[:, 1]) / (B2[:, 0] - B1[:, 0])
        tan_k = (kLine2 - kLine1) / (1 + kLine2 * kLine1)
        lines_arctan = math.atan(tan_k)
        lines_arctan = float('%.2f' % abs(-lines_arctan * 180.0 / 3.1415926))
        if save_dots:
            dbg_image_circles = process_dots.copy()
            dbg_image_circles = cv2.drawChessboardCorners(dbg_image_circles, (3, 3), corners, found)
            result_path = save_figure(dbg_image_circles, 'detection_dots', res_code)
            detection_dots_img_stream = return_img_stream(result_path)

        return lines_arctan, [original_img_stream, process_dots_img_stream, detection_dots_img_stream]
'''
采用传统方式实现图像增强后进行网线角度检测
'''
def dots_angle_measure_tr_sr(dots_region,save_dots=False,res_code=None):

    original_img_stream = None
    process_dots_img_stream = None
    detection_dots_img_stream = None
    lines_arctan = None
    if save_dots:
        result_path = save_figure(dots_region, 'original', res_code)
        original_img_stream = return_img_stream(result_path)
    # 锐化插值处理
    sigma_list = [15, 80, 200]
    process_dots = automatedMSRCR(
        dots_region,
        sigma_list
    )
    size = 4
    process_dots = cv2.resize(process_dots, None, fx=size, fy=size, interpolation=cv2.INTER_CUBIC)

    if save_dots:
        result_path = save_figure(process_dots, ' process_dots', res_code)
        process_dots_img_stream = return_img_stream(result_path)

    # 斑点检测
    found, corners = cv2.findCirclesGrid(process_dots, (3, 3), cv2.CALIB_CB_SYMMETRIC_GRID)

    if corners is None or corners.any() != None and corners.shape[0] < 2:
        # print("---------------------------------------------------------------")
        # print("corners status,", found)
        # print("---------------------------------------------------------------")
        return lines_arctan, [original_img_stream, process_dots_img_stream]

    elif corners.shape[0] >= 2:
        A1 = corners[0]
        A2 = corners[1]
        if (A1 == A2).all():
            # 斑点检测
            found, corners = cv2.findCirclesGrid(process_dots, (4, 4), cv2.CALIB_CB_SYMMETRIC_GRID)
            A1 = corners[0]
            A2 = corners[1]
            if (A1 == A2).all():
                return lines_arctan, [original_img_stream, process_dots_img_stream]
        B1 = corners[1]
        B2 = np.array([[0, corners[1][:, 1]]],dtype=object)

        kLine1 = (A2[:, 1] - A1[:, 1]) / (A2[:, 0] - A1[:, 0])
        kLine2 = (B2[:, 1] - B1[:, 1]) / (B2[:, 0] - B1[:, 0])
        tan_k = (kLine2 - kLine1) / (1 + kLine2 * kLine1)
        lines_arctan = math.atan(tan_k)
        lines_arctan = float('%.2f' % abs(-lines_arctan * 180.0 / 3.1415926))
        if save_dots:
            dbg_image_circles = process_dots.copy()
            dbg_image_circles = cv2.drawChessboardCorners(dbg_image_circles, (3, 3), corners, found)
            result_path = save_figure(dbg_image_circles, 'detection_dots', res_code)
            detection_dots_img_stream = return_img_stream(result_path)

        return lines_arctan, [original_img_stream, process_dots_img_stream, detection_dots_img_stream]