import cv2

import matplotlib.pyplot as plt

from libs.utils import *
from libs.modules import *

class Visualize_plt(object):

    def __init__(self, cfg=None):

        self.cfg = cfg
        self.height = cfg.height
        self.width = cfg.width
        self.size = to_tensor(np.float32(cfg.size))

        self.mean = np.array([cfg.mean], dtype=np.float32)
        self.std = np.array([cfg.std], dtype=np.float32)

        self.line = np.zeros((cfg.height, 3, 3), dtype=np.uint8)
        self.line[:, :, :] = 255

        self.param = {'linewidth': [0, 1, 2, 3, 4], 'color': ['yellow', 'red', 'lime']}

        self.show = {}


    def load_image(self, dir_name, file_name, name):
        # org image
        img = cv2.imread(dir_name + file_name)
        img = cv2.resize(img, (self.height, self.width))
        self.show[name] = img
        self.img = img

    def show_image(self):
        plt.figure()
        plt.imshow(self.img[:, :, [2, 1, 0]])

    def save_fig(self, path, name):
        mkdir(path)
        plt.axis('off')
        plt.savefig(path + name, bbox_inches='tight', pad_inches=0)
        plt.close()

    def save_img(self, dir_name, file_name, list):
        disp = self.line
        for i in range(len(list)):
            disp = np.concatenate((disp, self.show[list[i]], self.line), axis=1)

        mkdir(dir_name)
        cv2.imwrite(dir_name + file_name, disp)


    def draw_lines_plt(self, pts, idx1, idx2, linestyle='-', zorder=1):
        endpts = find_endpoints(pts, [self.width-1, self.height-1])

        pt_1 = (endpts[0], endpts[1])
        pt_2 = (endpts[2], endpts[3])
        plt.plot([pt_1[0], pt_2[0]], [pt_1[1], pt_2[1]],
                 linestyle=linestyle,
                 linewidth=self.param['linewidth'][idx1],
                 color=self.param['color'][idx2],
                 zorder=zorder)

    def draw_lines_cv(self, data, name, ref_name='img', color=(255, 0, 0), s=1):
        img = np.ascontiguousarray(np.copy(self.show[ref_name]))
        for i in range(data.shape[0]):
            pts = data[i]
            pt_1 = (pts[0], pts[1])
            pt_2 = (pts[2], pts[3])
            img = cv2.line(img, pt_1, pt_2, color, s)

        self.show[name] = img

    def display_for_train_detector(self, batch, out, idx):
        img_name = batch['img_name'][0]
        candidates = batch['train_data'][0].cuda()
        self.show['img_name'] = img_name

        img = to_np(batch['img_rgb'][0].permute(1, 2, 0))
        img = np.uint8(img[:, :, [2, 1, 0]] * 255)
        self.show['img'] = img

        pos_line = candidates[candidates[:, 5] == 1][:, :4]
        neg_line = candidates[candidates[:, 5] == 0][:, :4]

        cls_out = torch.argmax(out['cls'], dim=1)
        out_line = candidates[cls_out == 1][:, :4] + \
                   out['reg'][cls_out == 1] * self.size

        pos_line = to_np(pos_line)
        neg_line = to_np(neg_line)
        out_line = to_np2(out_line)

        self.draw_lines_cv(data=pos_line, name='pos')
        self.draw_lines_cv(data=neg_line, name='neg')
        self.draw_lines_cv(data=out_line, name='out')

        self.save_img(dir_name=self.cfg.output_dir + 'train/display/',
                      file_name=str(idx)+'.jpg',
                      list=['pos', 'neg', 'out'])

    def display_for_test(self, batch, out):
        img_name = batch['img_name'][0]
        self.load_image(dir_name=self.cfg.dataset_dir,
                        file_name=img_name,
                        name='img')
        self.show_image()

        if 'pri' in out.keys():
            pts_pri = to_np(out['pri'])
            self.draw_lines_plt(pts_pri[0], 4, 1, '--', zorder=2)
        if 'mul' in out.keys():
            pts_mul = to_np(out['mul'])
            for i in range(pts_mul.shape[0]):
                self.draw_lines_plt(pts_mul[i], 4, 0, '-', zorder=1)

        self.save_fig(path=self.cfg.output_dir + 'test/out_single/',
                      name=img_name.split('/')[-1][:-4] + '.png')

