import numpy as np
import sys
import os
from PIL import Image
from visu.helper_functions import save_image
from scipy.spatial.transform import Rotation as R
from helper import re_quat
import copy
import torch
import numpy as np
import k3d


class Visualizer():
    def __init__(self, p_visu, writer=None):
        if p_visu[-1] != '/':
            p_visu = p_visu + '/'
        self.p_visu = p_visu
        self.writer = writer

        if not os.path.exists(self.p_visu):
            os.makedirs(self.p_visu)

    def plot_estimated_pose(self, tag, epoch, img, points, trans=[[0, 0, 0]], rot_mat=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], cam_cx=0, cam_cy=0, cam_fx=0, cam_fy=0, store=False, jupyter=False, w=2):
        """
        tag := tensorboard tag 
        epoch := tensorboard epoche
        store := ture -> stores the image to standard path
        path := != None creats the path and store to it path/tag.png
        img:= original_image, [widht,height,RGB]
        points:= points of the object model [length,x,y,z]
        trans: [1,3]
        rot: [3,3]
        """

        img_d = copy.deepcopy(img)
        points = np.dot(points, rot_mat.T)
        points = np.add(points, trans[0, :])
        for i in range(0, points.shape[0]):
            p_x = points[i, 0]
            p_y = points[i, 1]
            p_z = points[i, 2]
            u = int(((p_x / p_z) * cam_fx) + cam_cx)
            v = int(((p_y / p_z) * cam_fy) + cam_cy)
            try:
                img_d[v - w:v + w + 1, u - w:u + w + 1, 0] = 0
                img_d[v - w:v + w + 1, u - w:u + w + 1, 1] = 255
                img_d[v - w:v + w + 1, u - w:u + w + 1, 0] = 0
            except:
                #print("out of bounce")
                pass

        if jupyter:
            display(Image.fromarray(img_d))

        if store:
            #store_ar = (img_d* 255).round().astype(np.uint8)
            #print("IMAGE D:" ,img_d,img_d.shape )
            save_image(img_d, tag=str(epoch) + tag, p_store=self.p_visu)
        if self.writer is not None:
            self.writer.add_image(tag, img_d.astype(
                np.uint8), global_step=epoch, dataformats='HWC')

    def plot_bounding_box(self, tag, epoch, img, rmin=0, rmax=0, cmin=0, cmax=0, str_width=2, store=False, jupyter=False, b=None):
        """
        tag := tensorboard tag 
        epoch := tensorboard epoche
        store := ture -> stores the image to standard path
        path := != None creats the path and store to it path/tag.png
        img:= original_image, [widht,height,RGB]

        """

        if isinstance(b, dict):
            rmin = b['rmin']
            rmax = b['rmax']
            cmin = b['cmin']
            cmax = b['cmax']

        # ToDo check Input data
        img_d = np.array(copy.deepcopy(img))

        c = [0, 0, 255]
        rmin_mi = max(0, rmin - str_width)
        rmin_ma = min(img_d.shape[0], rmin + str_width)

        rmax_mi = max(0, rmax - str_width)
        rmax_ma = min(img_d.shape[0], rmax + str_width)

        cmin_mi = max(0, cmin - str_width)
        cmin_ma = min(img_d.shape[1], cmin + str_width)

        cmax_mi = max(0, cmax - str_width)
        cmax_ma = min(img_d.shape[1], cmax + str_width)

        img_d[rmin_mi:rmin_ma, cmin:cmax, :] = c
        img_d[rmax_mi:rmax_ma, cmin:cmax, :] = c
        img_d[rmin:rmax, cmin_mi:cmin_ma, :] = c
        img_d[rmin:rmax, cmax_mi:cmax_ma, :] = c
        print("STORE", store)
        img_d = img_d.astype(np.uint8)
        if store:
            #store_ar = (img_d* 255).round().astype(np.uint8)
            save_image(img_d, tag=str(epoch) + tag, p_store=self.p_visu)
        if jupyter:
            display(Image.fromarray(img_d))
        if self.writer is not None:
            self.writer.add_image(tag, img_d.astype(
                np.uint8), global_step=epoch, dataformats='HWC')


def plot_pcd(x, point_size=0.005, c='g'):
    """
    x: point_nr,3
    """
    if c == 'b':
        k = 245
    elif c == 'g':
        k = 25811000
    elif c == 'r':
        k = 11801000
    elif c == 'black':
        k = 2580
    else:
        k = 2580
    colors = np.ones(x.shape[0]) * k
    plot = k3d.plot(name='points')
    plt_points = k3d.points(x, colors.astype(np.uint32), point_size=point_size)
    plot += plt_points
    plt_points.shader = '3d'
    plot.display()


def plot_two_pcd(x, y, point_size=0.005, c1='g', c2='r'):
    if c1 == 'b':
        k = 245
    elif c1 == 'g':
        k = 25811000
    elif c1 == 'r':
        k = 11801000
    elif c1 == 'black':
        k = 2580
    else:
        k = 2580

    if c2 == 'b':
        k2 = 245
    elif c2 == 'g':
        k2 = 25811000
    elif c2 == 'r':
        k2 = 11801000
    elif c2 == 'black':
        k2 = 2580
    else:
        k2 = 2580

    col1 = np.ones(x.shape[0]) * k
    col2 = np.ones(y.shape[0]) * k2
    plot = k3d.plot(name='points')
    plt_points = k3d.points(x, col1.astype(np.uint32), point_size=point_size)
    plot += plt_points
    plt_points = k3d.points(y, col2.astype(np.uint32), point_size=point_size)
    plot += plt_points
    plt_points.shader = '3d'
    plot.display()


class SequenceVisualizer():
    def __init__(self, seq_data, images_path, output_path=None):
        self.seq_data = seq_data
        self.images_path = images_path
        self.output_path = output_path

    def plot_points_on_image(self, seq_no, frame_no, jupyter=False, store=False, pose_type='filtered'):
        seq_data = self.seq_data
        images_path = self.images_path
        output_path = self.output_path
        frame = seq_data[seq_no][frame_no]
        unique_desig = frame['dl_dict']['unique_desig'][0]

        if pose_type == 'ground_truth':
            # ground truth
            t = frame['dl_dict']['gt_trans'].reshape(1, 3)
            rot_quat = re_quat(copy.deepcopy(
                frame['dl_dict']['gt_rot_wxyz'][0]), 'wxyz')
            rot = R.from_quat(rot_quat).as_matrix()
        elif pose_type == 'filtered':
            # filter pred
            t = np.array(frame['filter_pred']['t']).reshape(1, 3)
            rot_quat = re_quat(copy.deepcopy(
                frame['filter_pred']['r_wxyz']), 'wxyz')
            rot = R.from_quat(rot_quat).as_matrix()
        elif pose_type == 'final_pred_obs':
            # final pred
            t = np.array(frame['final_pred_obs']['t']).reshape(1, 3)
            rot_quat = re_quat(copy.deepcopy(
                frame['final_pred_obs']['r_wxyz']), 'wxyz')
            rot = R.from_quat(rot_quat).as_matrix()
        else:
            raise Exception('Pose type not implemented.')

        w = 2
        if type(unique_desig) != str:
            im = np.array(Image.open(
                images_path + unique_desig[0] + '-color.png'))  # ycb
        else:
            im = np.array(Image.open(
                images_path + unique_desig + '.png'))  # laval
        img_d = copy.deepcopy(im)

        dl_dict = frame['dl_dict']
        points = copy.deepcopy(
            seq_data[seq_no][0]['dl_dict']['model_points'][0, :, :])
        points = np.dot(points, rot.T)
        points = np.add(points, t[0, :])

        cam_cx = dl_dict['cam_cal'][0][0]
        cam_cy = dl_dict['cam_cal'][0][1]
        cam_fx = dl_dict['cam_cal'][0][2]
        cam_fy = dl_dict['cam_cal'][0][3]
        for i in range(0, points.shape[0]):
            p_x = points[i, 0]
            p_y = points[i, 1]
            p_z = points[i, 2]
            u = int(((p_x / p_z) * cam_fx) + cam_cx)
            v = int(((p_y / p_z) * cam_fy) + cam_cy)
            try:
                img_d[v - w:v + w + 1, u - w:u + w + 1, 0] = 0
                img_d[v - w:v + w + 1, u - w:u + w + 1, 1] = 255
                img_d[v - w:v + w + 1, u - w:u + w + 1, 0] = 0
            except:
                #print("out of bounds")
                pass

        img_disp = Image.fromarray(img_d)
        if jupyter:
            display(img_disp)

        if store:
            outpath = output_path + \
                '{}_{}_{}.png'.format(pose_type, seq_no, frame_no)
            img_disp.save(outpath, "PNG", compress_level=1)
            print("Saved image to {}".format(outpath))

    def save_sequence(self, seq_no, pose_type='filtered', name=''):
        for fn in range(len(self.seq_data)):
            self.plot_points_on_image(seq_no, fn, False, True, pose_type)
        if name:
            video_name = '{}_{}_{}'.format(name, pose_type, seq_no)
        else:
            video_name = '{}_{}'.format(pose_type, seq_no)
        cmd = "cd {} && ffmpeg -r 10 -i ./filtered_{}_%d.png -vcodec mpeg4 -y {}.mp4".format(
            self.output_path, seq_no, video_name)
        os.system(cmd)

[((127, 11, 127, 34), 'k3d.plot', 'k3d.plot', (), '', False, 'import k3d\n'), ((159, 11, 159, 34), 'k3d.plot', 'k3d.plot', (), '', False, 'import k3d\n'), ((36, 16, 36, 34), 'copy.deepcopy', 'copy.deepcopy', ({(36, 30, 36, 33): 'img'}, {}), '(img)', False, 'import copy\n'), ((37, 17, 37, 42), 'numpy.dot', 'np.dot', ({(37, 24, 37, 30): 'points', (37, 32, 37, 41): 'rot_mat.T'}, {}), '(points, rot_mat.T)', True, 'import numpy as np\n'), ((38, 17, 38, 44), 'numpy.add', 'np.add', ({(38, 24, 38, 30): 'points', (38, 32, 38, 43): 'trans[(0), :]'}, {}), '(points, trans[(0), :])', True, 'import numpy as np\n'), ((126, 13, 126, 32), 'numpy.ones', 'np.ones', ({(126, 21, 126, 31): 'x.shape[0]'}, {}), '(x.shape[0])', True, 'import numpy as np\n'), ((157, 11, 157, 30), 'numpy.ones', 'np.ones', ({(157, 19, 157, 29): 'x.shape[0]'}, {}), '(x.shape[0])', True, 'import numpy as np\n'), ((158, 11, 158, 30), 'numpy.ones', 'np.ones', ({(158, 19, 158, 29): 'y.shape[0]'}, {}), '(y.shape[0])', True, 'import numpy as np\n'), ((209, 16, 209, 33), 'copy.deepcopy', 'copy.deepcopy', ({(209, 30, 209, 32): 'im'}, {}), '(im)', False, 'import copy\n'), ((212, 17, 213, 68), 'copy.deepcopy', 'copy.deepcopy', ({(213, 12, 213, 67): "seq_data[seq_no][0]['dl_dict']['model_points'][(0), :, :]"}, {}), "(seq_data[seq_no][0]['dl_dict']['model_points'][(0), :, :])", False, 'import copy\n'), ((214, 17, 214, 38), 'numpy.dot', 'np.dot', ({(214, 24, 214, 30): 'points', (214, 32, 214, 37): 'rot.T'}, {}), '(points, rot.T)', True, 'import numpy as np\n'), ((215, 17, 215, 40), 'numpy.add', 'np.add', ({(215, 24, 215, 30): 'points', (215, 32, 215, 39): 't[(0), :]'}, {}), '(points, t[(0), :])', True, 'import numpy as np\n'), ((235, 19, 235, 41), 'PIL.Image.fromarray', 'Image.fromarray', ({(235, 35, 235, 40): 'img_d'}, {}), '(img_d)', False, 'from PIL import Image\n'), ((254, 8, 254, 22), 'os.system', 'os.system', ({(254, 18, 254, 21): 'cmd'}, {}), '(cmd)', False, 'import os\n'), ((21, 15, 21, 42), 'os.path.exists', 'os.path.exists', ({(21, 30, 21, 41): 'self.p_visu'}, {}), '(self.p_visu)', False, 'import os\n'), ((22, 12, 22, 36), 'os.makedirs', 'os.makedirs', ({(22, 24, 22, 35): 'self.p_visu'}, {}), '(self.p_visu)', False, 'import os\n'), ((81, 25, 81, 43), 'copy.deepcopy', 'copy.deepcopy', ({(81, 39, 81, 42): 'img'}, {}), '(img)', False, 'import copy\n'), ((54, 20, 54, 42), 'PIL.Image.fromarray', 'Image.fromarray', ({(54, 36, 54, 41): 'img_d'}, {}), '(img_d)', False, 'from PIL import Image\n'), ((106, 20, 106, 42), 'PIL.Image.fromarray', 'Image.fromarray', ({(106, 36, 106, 41): 'img_d'}, {}), '(img_d)', False, 'from PIL import Image\n'), ((184, 31, 185, 51), 'copy.deepcopy', 'copy.deepcopy', ({(185, 16, 185, 50): "frame['dl_dict']['gt_rot_wxyz'][0]"}, {}), "(frame['dl_dict']['gt_rot_wxyz'][0])", False, 'import copy\n'), ((204, 26, 205, 61), 'PIL.Image.open', 'Image.open', ({(205, 16, 205, 60): "images_path + unique_desig[0] + '-color.png'"}, {}), "(images_path + unique_desig[0] + '-color.png')", False, 'from PIL import Image\n'), ((207, 26, 208, 52), 'PIL.Image.open', 'Image.open', ({(208, 16, 208, 51): "images_path + unique_desig + '.png'"}, {}), "(images_path + unique_desig + '.png')", False, 'from PIL import Image\n'), ((186, 18, 186, 39), 'scipy.spatial.transform.Rotation.from_quat', 'R.from_quat', ({(186, 30, 186, 38): 'rot_quat'}, {}), '(rot_quat)', True, 'from scipy.spatial.transform import Rotation as R\n'), ((190, 31, 191, 47), 'copy.deepcopy', 'copy.deepcopy', ({(191, 16, 191, 46): "frame['filter_pred']['r_wxyz']"}, {}), "(frame['filter_pred']['r_wxyz'])", False, 'import copy\n'), ((189, 16, 189, 51), 'numpy.array', 'np.array', ({(189, 25, 189, 50): "frame['filter_pred']['t']"}, {}), "(frame['filter_pred']['t'])", True, 'import numpy as np\n'), ((192, 18, 192, 39), 'scipy.spatial.transform.Rotation.from_quat', 'R.from_quat', ({(192, 30, 192, 38): 'rot_quat'}, {}), '(rot_quat)', True, 'from scipy.spatial.transform import Rotation as R\n'), ((196, 31, 197, 50), 'copy.deepcopy', 'copy.deepcopy', ({(197, 16, 197, 49): "frame['final_pred_obs']['r_wxyz']"}, {}), "(frame['final_pred_obs']['r_wxyz'])", False, 'import copy\n'), ((195, 16, 195, 54), 'numpy.array', 'np.array', ({(195, 25, 195, 53): "frame['final_pred_obs']['t']"}, {}), "(frame['final_pred_obs']['t'])", True, 'import numpy as np\n'), ((198, 18, 198, 39), 'scipy.spatial.transform.Rotation.from_quat', 'R.from_quat', ({(198, 30, 198, 38): 'rot_quat'}, {}), '(rot_quat)', True, 'from scipy.spatial.transform import Rotation as R\n')]