import json
import os

import torch
import torch.nn.functional as F
import tritonclient.http as httpclient
import requests
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm

from app.service.outfit_matcher.dataset import FashionDataset
from app.service.outfit_matcher.foco import extract_main_colors
from app.service.outfit_matcher.outfit_evaluator import evaluate_outfits, visualize


class OutfitMatcherHon:
    def __init__(self, outfits):
        self.outfits = outfits
        self.tritonclient = httpclient.InferenceServerClient(url="10.1.1.240:10010")

    @staticmethod
    def imnormalize(img, mean, std, to_rgb=True):
        """Normalize an image with mean and std.

        Args:
            img (ndarray): Image to be normalized.
            mean (ndarray): The mean to be used for normalize.
            std (ndarray): The std to be used for normalize.
            to_rgb (bool): Whether to convert to rgb.

        Returns:
            ndarray: The normalized image.
        """
        img = img.copy().astype(np.float32)
        assert img.dtype != np.uint8
        mean = np.float64(mean.reshape(1, -1))
        stdinv = 1 / np.float64(std.reshape(1, -1))
        if to_rgb:
            cv2.cvtColor(img, cv2.COLOR_BGR2RGB, img)  # inplace
        cv2.subtract(img, mean, img)  # inplace
        cv2.multiply(img, stdinv, img)  # inplace
        return img

    @staticmethod
    def load_image(img_path):
        if 'http' in img_path:
            file = requests.get(img_path)
            image = cv2.imdecode(np.fromstring(file.content, np.uint8), 1)
            image = Image.fromarray(image.astype('uint8'), 'RGB')
        else:
            image = Image.open(img_path).convert('RGB')
        return np.array(image)

    @staticmethod
    def resize_image(img):
        """
        Args:
            img: ndarray (height, width, channel)
        """
        resized_img = cv2.resize(img, (224, 224), dst=None, interpolation=1)
        return resized_img

    @staticmethod
    def pad_array(input_value):
        """pad List of Array into same batch size

        Args:
            input_value: List of numpy arrary need to be padded

        Returns:
            Tensor: [batch_dim, max_dim, original_tensor_size]
        """
        max_dim = max([len(x) for x in input_value])
        mask = np.zeros((len(input_value), max_dim), dtype=np.float32)

        # Pad each array
        padded_arrays = []
        for i, array in enumerate(input_value):
            # Compute padding amount along the pad dimension
            pad_dim = max_dim - array.shape[0]
            consistent_shape = array.shape[1:]
            pad_widths = [(0, pad_dim)] + [(0, 0)] * len(consistent_shape)
            padded_array = np.pad(array, pad_widths, mode='constant', constant_values=0)
            padded_arrays.append(padded_array)

            mask[i, array.shape[0]:] = float("-inf")

        # Stack the padded arrays and change the dimension
        batched_arrays = np.stack(padded_arrays, axis=0)
        return batched_arrays, mask

    def preprocess(self):
        outfit_images = []
        outfit_colors = []
        for outfit in self.outfits:
            images = []
            colors = []
            for item in outfit:
                image = self.load_image(item["image_path"])
                image = self.resize_image(image)
                normalized_image = self.imnormalize(image,
                                                    mean=np.array([208.32996145, 201.28227452, 198.47047691], dtype=np.float32),
                                                    std=np.array([75.48939648, 80.47423057, 82.21144189], dtype=np.float32))
                images.append(normalized_image.transpose(2, 0, 1))
                color = extract_main_colors(image)
                colors.append(color)
            images = np.stack(images, axis=0)
            outfit_images.append(images)  # List[(items, 3, 224, 224)]
            colors = np.stack(colors, axis=0)
            outfit_colors.append(colors)
        outfit_images, mask = self.pad_array(outfit_images)
        outfit_colors, _ = self.pad_array(outfit_colors)
        return outfit_images, outfit_colors, mask

    def get_result(self):
        # start = time.time()
        image, color, mask = self.preprocess()
        # print(start - time.time())
        # transformed_img = image.astype(np.float32)
        # 输入集
        inputs = [
            httpclient.InferInput("input__0", image.shape, datatype="FP32"),
            httpclient.InferInput("input__1", color.shape, datatype="FP32"),
            httpclient.InferInput("input__2", mask.shape, datatype="FP32"),
        ]
        inputs[0].set_data_from_numpy(image.astype(np.float32), binary_data=True)
        inputs[1].set_data_from_numpy(color.astype(np.float32), binary_data=True)
        inputs[2].set_data_from_numpy(mask.astype(np.float32), binary_data=True)
        # 输出集
        outputs = [
            httpclient.InferRequestedOutput("output__0", binary_data=True),
        ]
        results = self.tritonclient.infer(model_name="outfit_matcher_hon", inputs=inputs, outputs=outputs)
        # 推理
        # 取结果
        inference_output1 = torch.from_numpy(results.as_numpy("output__0"))
        return inference_output1  # Shape (N, 1)


if __name__ == '__main__':
    with open("./test_param/recommendation_test.json", "r") as f:
        param = json.load(f)
    fashion_dataset = FashionDataset(param["database"])
    for item in tqdm(param["query"]):
        outfits = fashion_dataset.generate_outfit(item, param["topk"], param["max_outfits"])
        service = OutfitMatcherHon(outfits=outfits)
        scores = service.get_result()
        visualize(outfits, scores, param["topk"], best=True,
                  output_path=os.path.join(
                      r"E:\workspace\outfit_matcher\2024 SS Outfit",
                      f"{item['item_name']}_best_{param['topk']}.png"
                  ))
        visualize(outfits, scores, param["topk"], best=False,
                  output_path=os.path.join(
                      r"E:\workspace\outfit_matcher\2024 SS Outfit",
                      f"{item['item_name']}_worst_{param['topk']}.png"
                  ))
        a = 1