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])