"""Inferrer"""
from PIL import Image

import torch
import numpy as np
import matplotlib.pyplot as plt 
import cv2

from utils.load import load_yaml
from model import get_model
from dataloader.transform import DataTransform


class Inferrer():
    """SSDでの予測と画像の表示をまとめて行うクラス"""

    def __init__(self, configfile):
        # Config
        config = load_yaml(configfile)

        self.model = get_model(config, is_eval=True)
        self.model.build(is_eval=True)
        self.net = self.model.model

        # 重みの読み込み
        self.net_weights = torch.load(config['infer']['weight_path'], map_location={'cuda:0': 'cpu'})
        self.net.load_state_dict(self.net_weights, strict=False)

        self.classes = self.model.classes
        self.data_confidence_level = config['infer']['data_confidence_level']
        self.color_mean = config['data']['color_mean']  # (BGR)の色の平均値
        self.input_size = config['data']['input_size']  # 画像のinputサイズを300×300にする
        self.transform = DataTransform(self.input_size, self.color_mean)  # 前処理クラス

    def show(self, image_file_path):
        """
        物体検出の予測結果を表示をする関数。
        Parameters
        ----------
        image_file_path:  str
            画像のファイルパス
        data_confidence_level: float
            予測で発見とする確信度の閾値
        Returns
        -------
        なし。rgb_imgに物体検出結果が加わった画像が表示される。
        """

        img = cv2.imread(image_file_path)  # [高さ][幅][色BGR]
        input_height, input_width, _ = img.shape  # 画像のサイズを取得

        rgb_img, predict_bbox, pre_dict_label_index, scores = self.ssd_predict(
            image_file_path, self.data_confidence_level)

        img = self.vis_bbox(rgb_img, bbox=predict_bbox, label_index=pre_dict_label_index,
                      scores=scores, label_names=self.classes,
                      crop_height=input_height, crop_width=input_width)
        return img

    def ssd_predict(self, image_file_path, data_confidence_level=0.5):
        """
        SSDで予測させる関数。
        Parameters
        ----------
        image_file_path:  strt
            画像のファイルパス
        dataconfidence_level: float
            予測で発見とする確信度の閾値
        Returns
        -------
        rgb_img, true_bbox, true_label_index, predict_bbox, pre_dict_label_index, scores
        """

        # rgbの画像データを取得
        img = cv2.imread(image_file_path)  # [高さ][幅][色BGR]
        height, width, channels = img.shape  # 画像のサイズを取得
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # 画像の前処理
        phase = "eval"
        img_transformed, boxes, labels = self.transform(
            img, phase, "", "")  # アノテーションが存在しないので""にする。
        img = torch.from_numpy(
            img_transformed[:, :, (2, 1, 0)]).permute(2, 0, 1)

        # SSDで予測
        self.net.eval()  # ネットワークを推論モードへ
        x = img.unsqueeze(0)  # ミニバッチ化：torch.Size([1, 3, 300, 300])

        detections = self.net(x)
        # detectionsの形は、torch.Size([1, 21, 200, 5])  ※200はtop_kの値

        # confidence_levelが基準以上を取り出す
        predict_bbox = []
        pre_dict_label_index = []
        scores = []
        detections = detections.cpu().detach().numpy()

        # 条件以上の値を抽出
        find_index = np.where(detections[:, 0:, :, 0] >= data_confidence_level)
        detections = detections[find_index]
        for i in range(len(find_index[1])):  # 抽出した物体数分ループを回す
            if (find_index[1][i]) > 0:  # 背景クラスでないもの
                sc = detections[i][0]  # 確信度
                bbox = detections[i][1:] * [width, height, width, height]
                # find_indexはミニバッチ数、クラス、topのtuple
                lable_ind = find_index[1][i]-1
                # （注釈）
                # 背景クラスが0なので1を引く

                # 返り値のリストに追加
                predict_bbox.append(bbox)
                pre_dict_label_index.append(lable_ind)
                scores.append(sc)

        return rgb_img, predict_bbox, pre_dict_label_index, scores

    def vis_bbox(self, rgb_img, bbox, label_index, scores, label_names, crop_height, crop_width,):
        """
        物体検出の予測結果を画像で表示させる関数。
        Parameters
        ----------
        rgb_img:rgbの画像
            対象の画像データ
        bbox: list
            物体のBBoxのリスト
        label_index: list
            物体のラベルへのインデックス
        scores: list
            物体の確信度。
        label_names: list
            ラベル名の配列
        Returns
        -------
        なし。rgb_imgに物体検出結果が加わった画像が表示される。
        """

        # 枠の色の設定
        num_classes = len(label_names)  # クラス数（背景のぞく）
        colors = plt.cm.hsv(np.linspace(0, 1, num_classes)).tolist()

        # 画像の表示
        fig = plt.figure(figsize=(crop_width/100, crop_height/100))
        plt.imshow(rgb_img)
        plt.axis("off")
        currentAxis = plt.gca()

        # BBox分のループ
        for i, bb in enumerate(bbox):

            # ラベル名
            label_name = label_names[label_index[i]]
            color = colors[label_index[i]]  # クラスごとに別の色の枠を与える

            # 枠につけるラベル 例：person;0.72
            if scores is not None:
                sc = scores[i]
                display_txt = '%s: %.2f' % (label_name, sc)
            else:
                display_txt = '%s: ans' % (label_name)

            # 枠の座標
            xy = (bb[0], bb[1])
            width = bb[2] - bb[0]
            height = bb[3] - bb[1]

            # 長方形を描画する
            currentAxis.add_patch(plt.Rectangle(
                xy, width, height, fill=False, edgecolor=color, linewidth=2))

            # 長方形の枠の左上にラベルを描画する
            currentAxis.text(xy[0], xy[1], display_txt, bbox={
                             'facecolor': color, 'alpha': 0.5})
                             
        fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
        fig.canvas.draw()
        im = np.array(fig.canvas.renderer.buffer_rgba())
        # im = np.array(fig.canvas.renderer._renderer) # matplotlibが3.1より前の場合

        img = Image.fromarray(im)
        img = img.convert('RGB')
        # 元の画像サイズにセンタークロップ
        img_width, img_height = img.size
        # img = img.crop(((img_width - crop_width) // 2,
        #                  (img_height - crop_height) // 2,
        #                  (img_width + crop_width) // 2,
        #                  (img_height + crop_height) // 2))

        return img
        # img.save('./test_output.png', quality=95)
