#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：trinity_client
@File    ：server.py
@Author  ：周成融
@Date    ：2023/9/16 18:31:08
@detail  ：
"""
from skimage import transform
import cv2
import mmcv
import numpy as np
import pandas as pd
from minio import Minio
import tritonclient.http as httpclient
import torch

from app.core.config import *
from app.schemas.attribute_retrieve import CategoryRecognitionModel
from app.service.utils.oss_client import oss_get_image


class CategoryRecognition:
    def __init__(self, request_data):
        self.attr_type = pd.read_csv(CATEGORY_PATH)
        # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
        self.request_data = []
        self.triton_client = httpclient.InferenceServerClient(url=DESIGN_MODEL_URL)
        for sketch in request_data:
            self.request_data.append(
                {
                    'obj': self.preprocess(self.get_image(sketch.sketch_img_url)),
                    'colony': sketch.colony,
                    'sketch_img_url': sketch.sketch_img_url,
                }
            )

    @staticmethod
    def preprocess(img):
        img = mmcv.imread(img)
        # ori_shape = img.shape[:2]
        img_scale = (224, 224)
        img = cv2.resize(img, img_scale)
        img = mmcv.imnormalize(img, mean=np.array([123.675, 116.28, 103.53]), std=np.array([58.395, 57.12, 57.375]), to_rgb=True)
        preprocessed_img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
        return preprocessed_img

    def get_image(self, url):
        # Get data of an object.
        # Read data from response.
        # response = self.minio_client.get_object(url.split("/", 1)[0], url.split("/", 1)[1])
        # img = np.frombuffer(response.data, np.uint8)  # 转成8位无符号整型
        # img = cv2.imdecode(img, cv2.IMREAD_COLOR)  # 解码
        img = oss_get_image(bucket=url.split("/", 1)[0], object_name=url.split("/", 1)[1], data_type="cv2")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img

    def get_category(self, image):
        inputs = [
            httpclient.InferInput("input__0", image.shape, datatype="FP32")
        ]
        inputs[0].set_data_from_numpy(image, binary_data=True)
        results = self.triton_client.infer(model_name="attr_retrieve_category", inputs=inputs)
        inference_output = torch.from_numpy(results.as_numpy(f'output__0'))

        scores = inference_output.detach().numpy()

        colattr = list(self.attr_type['labelName'])

        task = self.attr_type['taskName'][0]

        maxsc = np.max(scores[0][:5])
        indexs = np.argwhere(scores == maxsc)[:, 1]

        # attr = {}
        # attr[task] = []
        # for i in range(len(indexs)):
        #     atr = colattr[indexs[i]]
        #     attr[task].append(atr)

        return colattr[indexs[0]]

    def get_result(self):
        for sketch in self.request_data:
            if sketch['colony'] == "Female":
                sketch['category'] = self.get_category(sketch["obj"])
            elif sketch['colony'] == "Male":
                category = self.get_category(sketch["obj"])
                if category == 'Trousers' or category == 'Skirt':
                    sketch['category'] = 'Bottoms'
                elif category == 'Blouse' or category == 'Dress':
                    sketch['category'] = 'Tops'
                else:
                    sketch['category'] = 'Outwear'
            # 删除图片对象
            del sketch["obj"]
        return self.request_data


class Rescale(object):
    """Rescale the image in a sample to a given size.

    Args:
        output_size (tuple or int): Desired output size. If tuple, output is
            matched to output_size. If int, smaller of image edges is matched
            to output_size keeping aspect ratio the same.
    """

    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        self.output_size = output_size

    def __call__(self, image, landmarks):
        h, w = image.shape[:2]
        if isinstance(self.output_size, int):
            if h > w:
                new_h, new_w = self.output_size * h / w, self.output_size
            else:
                new_h, new_w = self.output_size, self.output_size * w / h
        else:
            new_h, new_w = self.output_size

        new_h, new_w = int(new_h), int(new_w)
        img = transform.resize(image, (new_h, new_w), mode='constant')
        # h and w are swapped for landmarks because for images,
        # x and y axes are axis 1 and 0 respectively
        landmarks = landmarks * [new_w / w, new_h / h]
        return img, landmarks


if __name__ == '__main__':
    data = [
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/ae976103-d7ec-4eed-b5d1-3e5f04d8be26.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/6d7d97a7-5a7d-48bd-9e14-b51119b48620.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/f2437141-1104-40a5-bcb9-f436088698bb.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/07af8613-eb2e-44fd-97cb-a97249a5754c.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/bac9fb15-6860-4112-ac97-f0dea079da75.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/11d59844-effa-4590-82f9-9ea382c76126.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/849bf94c-66b8-42f5-8c2e-c1c1f4c8d0e0.jpg",
            "colony": "Female"
        },
        {
            "sketch_img_url": "aida-users/89/sketchboard/female/Dress/e6724ab7-8d3f-4677-abe0-c3e42ab7af85.jpeg",
            "colony": "Female"
        }
    ]
    rq_data = [CategoryRecognitionModel(sketch_img_url=d['sketch_img_url'], colony=d['colony']) for d in data]
    server = CategoryRecognition(rq_data)

    print(server.get_result())
