import subprocess
subprocess.call("mim install 'mmengine>=0.6.0'", shell=True)
subprocess.call("mim install 'mmcv>=2.0.0rc4,<2.1.0'", shell=True)
subprocess.call("min install 'mmdet>=3.0.0,<4.0.0'", shell=True)
subprocess.call("mim install 'mmyolo'", shell=True)
subprocess.call("mim install 'mmpose'", shell=True)
subprocess.call("mim install 'mmpretrain'", shell=True)

import numpy as np
import gradio as gr
import requests
import base64
import pandas as pd
import cv2
from typing import Tuple
from PIL import Image
from io import BytesIO
from skimage import color
import os

from Model.Model6.model6_inference import main as model6_inferencer
from mmyolo.utils import register_all_modules

register_all_modules()


def get_access_token(refatch=False) -> str:
    """获取百度AI的access_token
    :param refatch:是否重新获取access_token
    :return:返回access_token"""
    if refatch:
        # client_id 为官网获取的AK， client_secret 为官网获取的SK
        client_id = '7OtH60uo01ZNYN4yPyahlRSx'
        client_secret = 'D5AxcUpyQyIA7KgPplp7dnz5tM0UIljy'
        host = 'https://aip.baidubce.com/oauth/2.0/token?' \
               'grant_type=client_credentials&client_id=%s&client_secret=%s' % (client_id, client_secret)
        response = requests.get(host)
        # print(response)
        if response:
            return response.json()['access_token']
    else:
        r"""
        {"refresh_token":"25.24b9368ce91f9bd62c8dad38b3436800.315360000.2007815067.282335-30479502",
         "expires_in":2592000,
         "session_key":
            "9mzdWT\/YmQ7oEi9WCRWbXd0YCcrSYQY6kKZjObKunlcKcZt95j9\/q1aJqbVXihpQOXK84o5WLJ8e7d4cXOi0VUJJcz5YEQ==",
         "access_token":"24.becefee37aba38ea43c546fc154d3016.2592000.1695047067.282335-30479502",
         "scope":"public brain_all_scope brain_body_analysis brain_body_attr brain_body_number brain_driver_behavior 
                  brain_body_seg brain_gesture_detect brain_body_tracking brain_hand_analysis wise_adapt 
                  lebo_resource_base lightservice_public hetu_basic lightcms_map_poi kaidian_kaidian 
                  ApsMisTest_Test\u6743\u9650 vis-classify_flower lpq_\u5f00\u653e cop_helloScope 
                  ApsMis_fangdi_permission smartapp_snsapi_base smartapp_mapp_dev_manage iop_autocar oauth_tp_app 
                  smartapp_smart_game_openapi oauth_sessionkey smartapp_swanid_verify smartapp_opensource_openapi 
                  smartapp_opensource_recapi fake_face_detect_\u5f00\u653eScope 
                  vis-ocr_\u865a\u62df\u4eba\u7269\u52a9\u7406 idl-video_\u865a\u62df\u4eba\u7269\u52a9\u7406 
                  smartapp_component smartapp_search_plugin avatar_video_test b2b_tp_openapi b2b_tp_openapi_online 
                  smartapp_gov_aladin_to_xcx","session_secret":"5c8c3dbb80b04f58bb33aa8077758679"
        }
        """
        access_token = "24.becefee37aba38ea43c546fc154d3016.2592000.1695047067.282335-30479502"
        return access_token


def resize_image(img, max_length=2048, min_length=50) -> Tuple[np.ndarray, bool]:
    """Ensure that the longest side is shorter than 2048px and the shortest side is longer than 50px.
    :param img: 前端传入的图片
    :param max_length: 最长边像素
    :param min_length: 最短边像素
    :return: 返回处理后的图片和是否进行了resize的标志
    """
    flag = False
    max_side = max(img.shape[0], img.shape[1])
    min_side = min(img.shape[0], img.shape[1])
    if max_side > max_length:
        scale = max_length / max_side
        img = cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale)))
        flag = True
    if min_side < min_length:
        scale = min_length / min_side
        img = cv2.resize(img, (int(img.shape[1] * scale), int(img.shape[0] * scale)))
        flag = True
    return img, flag


def model1_det(x):
    """人体检测与属性识别
    :param x:前端传入的图片
    :return:返回检测结果
    """

    def _Baidu_det(img):
        """调用百度AI接口进行人体检测与属性识别
        :param img:前端传入的图片，格式为numpy.ndarray
        :return:返回检测结果
        """
        request_url = "https://aip.baidubce.com/rest/2.0/image-classify/v1/body_attr"
        # 保存图片到本地
        cv2.imwrite('test.jpg', img)
        # 二进制方式打开图片文件
        f = open('test.jpg', 'rb')
        hex_image = base64.b64encode(f.read())
        # 选择二进制图片和需要输出的属性（12个）
        params = {
            "image": hex_image,
            "type": "gender,age,upper_wear,lower_wear,upper_color,lower_color,"
                    "orientation,upper_cut,lower_cut,side_cut,occlusion,is_human"
        }
        access_token = get_access_token()
        request_url = request_url + "?access_token=" + access_token
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        if response:
            return response.json()

    def _get_attributes_list(r) -> dict:
        """获取人体属性列表
        :param r:百度AI接口返回的json数据
        :return:返回人体属性列表
        """
        all_humans_attributes_list = {}
        person_num = r['person_num']
        print('person_num:', person_num)
        for human_idx in range(person_num):
            attributes_dict = r['person_info'][human_idx]['attributes']
            attributes_list = []
            for key, value in attributes_dict.items():
                attribute = [key, value['name'], value['score']]
                attributes_list.append(attribute)
            new_value = ['attribute', 'attribute_value', 'accuracy']
            attributes_list.insert(0, new_value)
            df = pd.DataFrame(attributes_list[1:], columns=attributes_list[0])
            all_humans_attributes_list[human_idx] = df
        return all_humans_attributes_list

    def _show_img(img, bboxes):
        """显示图片
        :param img:前端传入的图片
        :param bboxes:检测框坐标
        :return:处理完成的图片 """
        line_width = int(max(img.shape[1], img.shape[0]) / 400)
        for bbox in bboxes:
            left, top, width, height = bbox['left'], bbox['top'], bbox['width'], bbox['height']
            right, bottom = left + width, top + height
            for i in range(left, right):
                img[top:top + line_width, i] = [255, 0, 0]
                img[bottom - line_width:bottom, i] = [255, 0, 0]
            for i in range(top, bottom):
                img[i, left:left + line_width] = [255, 0, 0]
                img[i, right - line_width:right] = [255, 0, 0]
        return img

    result = _Baidu_det(x)
    HAs_list = _get_attributes_list(result)
    locations = []
    for i in range(len(result['person_info'])):
        locations.append(result['person_info'][i]['location'])

    return _show_img(x, locations), f"模型检测到的人数为：{result['person_num']}人"


def model2_rem(x):
    """背景消除
    :param x: 前端传入的图片
    :return: 返回处理后的图片
    """

    def _Baidu_rem(img):
        """调用百度AI接口进行背景消除
        :param img: 前端传入的图片，格式为numpy.ndarray
        :return: 返回处理后的图片
        """
        request_url = "https://aip.baidubce.com/rest/2.0/image-classify/v1/body_seg"
        bgr_image = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        cv2.imwrite('test.jpg', bgr_image)
        f = open('test.jpg', 'rb')
        hex_image = base64.b64encode(f.read())
        params = {"image": hex_image}
        access_token = get_access_token()
        request_url = request_url + "?access_token=" + access_token
        headers = {'content-type': 'application/x-www-form-urlencoded'}
        response = requests.post(request_url, data=params, headers=headers)
        if response:
            encoded_image = response.json()["foreground"]
            decoded_image = base64.b64decode(encoded_image)
            image = Image.open(BytesIO(decoded_image))
            image_array = np.array(image)
            return image_array

    resized_x, resized_f = resize_image(x)
    new_img = _Baidu_rem(resized_x)
    if resized_f:
        resized_f = "图片尺寸已被修改至合适大小"
    else:
        resized_f = "图片尺寸无需修改"

    return new_img, resized_f


def model3_ext(x: np.ndarray, num_clusters=12):
    """主色调提取
    :param x: 前端传入的图片
    :param num_clusters: 聚类的数量
    :return: 返回处理后的图片,是否进行了resize的标志,颜色列表"""

    def _find_name(r, g, b):
        """根据颜色值查找颜色名称。
        :param r: 红色值
        :param g: 绿色值
        :param b: 蓝色值
        :return:返回颜色名称
        """
        # turn RGB to Lab
        lab = color.rgb2lab([[r / 255, g / 255, b / 255]])[0]
        for i in range(len(df)):
            # culcuate the minimum chromatic distance
            df['distance'] = np.sqrt((df['L'] - lab[0]) ** 2 + (df['a'] - lab[1]) ** 2 + (df['b'] - lab[2]) ** 2)
        # find the color name, whose chromatic distance is the minimum, and the corresponding distance
        name = df[df['distance'] == df['distance'].min()]['name'].values[0]
        distance = df[df['distance'] == df['distance'].min()]['distance'].values[0]

        return name, distance

    def _cluster(img, NUM_CLUSTERS):
        """K-means 聚类提取主色调
        :param img: 前端传入的图片
        :param NUM_CLUSTERS: 聚类的数量
        :return: 返回聚类结果
        """
        h, w, ch = img.shape
        reshaped_x = np.float32(img.reshape((-1, 4)))
        new_data_list = []
        for i in range(len(reshaped_x)):
            if reshaped_x[i][3] < 100:
                continue
            else:
                new_data_list.append(reshaped_x[i])
        reshaped_x = np.array(new_data_list)
        reshaped_x = np.delete(reshaped_x, 3, axis=1)
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
        NUM_CLUSTERS = NUM_CLUSTERS
        ret, label, center = cv2.kmeans(reshaped_x, NUM_CLUSTERS, None, criteria,
                                        NUM_CLUSTERS, cv2.KMEANS_RANDOM_CENTERS)
        clusters = np.zeros([NUM_CLUSTERS], dtype=np.int32)
        for i in range(len(label)):
            clusters[label[i][0]] += 1
        clusters = np.float32(clusters) / float(len(reshaped_x))
        center = np.int32(center)
        x_offset = 0
        card = np.zeros((50, w, 3), dtype=np.uint8)
        color_list = []
        for c in np.argsort(clusters)[::-1]:
            dx = int(clusters[c] * w)
            r = center[c][0]
            g = center[c][1]
            b = center[c][2]
            cv2.rectangle(card, (x_offset, 0), (x_offset + dx, 50),
                          (int(r), int(g), int(b)), -1)
            color_list.append([r, g, b, str(round(clusters[c]*100, 2)) + '%'])
            x_offset += dx

        return card, resized_f, color_list

    file = '中国传统色_集合.xlsx'
    df = pd.read_excel(file, sheet_name='Sheet1')[['name', 'L', 'a', 'b']]

    resized_x, resized_f = resize_image(x)
    card, resized_f, c_list = _cluster(resized_x, num_clusters)

    for c in c_list:
        c_name, c_distance = _find_name(c[0], c[1], c[2])
        c.append(c_name)
        c.append(c_distance)

    if resized_f:
        resized_f = "图片尺寸已被修改至合适大小"
    else:
        resized_f = "图片尺寸无需修改"

    c_df = pd.DataFrame(c_list, columns=['R', 'G', 'B', '比例', '颜色名称', '色差ΔE'])

    return card, resized_f, c_df


def model4_clo(x_path: str):
    def _get_result(input_path: str, cls_results: dict) -> pd.DataFrame:
        """convert the results of model6_2 to a dataframe
        :param input_path: the (absolute) path of the image
        :param cls_results: the results of model6_2

        :return: a dataframe to display on the web
        """
        result_pd = []
        img_name = os.path.basename(input_path)
        pred_profile = cls_results[img_name][0]['pred_class']
        pred_score = round(cls_results[img_name][0]['pred_score'], 2)
        result_pd.append([img_name, pred_profile, pred_score])
        df = pd.DataFrame(result_pd, columns=None)
        return df

    output_path_root = 'upload_to_web_tmp'
    if not os.path.exists(output_path_root):
        os.mkdir(output_path_root)
    cls_result = model6_inferencer(x_path, output_path_root)

    if cls_result:
        # use np to read image·
        x_name = os.path.basename(x_path)
        pred_x = np.array(Image.open(os.path.join(output_path_root, 'visualizations', x_name)))

        return pred_x, _get_result(x_path, cls_result), "识别成功!"
    # TODO: 完善识别失败时的处理（model6_inference.py中）[important]
    return x_path, pd.DataFrame(), "未检测到服装"


with gr.Blocks() as demo:
    gr.Markdown("# 服装图像识别模块——功能演示")
    with gr.Tab("人体检测模型"):
        with gr.Row():
            model1_input = gr.Image(height=400)
            model1_output_img = gr.Image(height=400)
            # model1_output_df = gr.DataFrame()
        model1_button = gr.Button("开始检测")
    with gr.Tab("背景消除模型"):
        with gr.Row():
            model2_input = gr.Image(height=400)
            model2_output_img = gr.Image(height=400)
        model2_button = gr.Button("开始消除")
    with gr.Tab('主色调提取'):
        with gr.Row():
            with gr.Column():
                model3_input = gr.Image(height=400, image_mode='RGBA')
                model3_slider = gr.Slider(minimum=1, maximum=20, step=1, value=12,
                                          min_width=400, label="聚类数量")
            with gr.Column():
                model3_output_img = gr.Image(height=100)
                model3_output_df = gr.DataFrame(headers=['R', 'G', 'B', '比例', '颜色名称', '色差ΔE'],
                                                datatype=['number', 'number', 'number', 'str', 'str', 'number'])

        model3_button = gr.Button("开始提取")
    with gr.Tab("廓形识别"):
        with gr.Row():
            model4_input = gr.Image(height=400, type="filepath")
            model4_output_img = gr.Image(height=400)
            model4_output_df = gr.DataFrame(headers=['img_name', 'pred_profile', 'pred_score'],
                                            datatype=['str', 'str', 'number'])
        model4_button = gr.Button("开始识别")
    # 设置折叠内容
    with gr.Accordion("模型运行信息"):
        running_info = gr.Markdown("等待输入和运行...")

    model1_button.click(model1_det, inputs=model1_input, outputs=[model1_output_img, running_info])
    model2_button.click(model2_rem, inputs=model2_input, outputs=[model2_output_img, running_info])
    model3_button.click(model3_ext,
                        inputs=[model3_input, model3_slider],
                        outputs=[model3_output_img, running_info, model3_output_df])
    model4_button.click(model4_clo, inputs=model4_input, outputs=[model4_output_img, model4_output_df, running_info])
demo.launch()