import os

import tensorflow as tf
from tqdm import tqdm
import numpy as np
from data.waymo_generator import Dataset
# tf.enable_eager_execution()
import train.waymo.waymo_config as config
from models.tf_ops.loader.sampling import grid_sampling
from models.tf_ops.loader.bbox_utils import get_roi_bbox
from models.tf_ops.loader.pooling import la_roi_pooling_fast
from models.tf_ops.loader.others import roi_filter, voxel2col
from models.tf_ops.test.test_utils import get_rgbs_from_coors, plot_points_from_roi_voxels, fetch_instance, plot_points

os.environ["CUDA_VISIBLE_DEVICES"] = "0"



anchor_size = [1.6, 3.9, 1.5]
batch_size = 4
epoch = 1


if __name__ == '__main__':
    WaymoDataset = Dataset(task='train',
                           batch_size=batch_size,
                           config=config.aug_config,
                           num_worker=6,
                           hvd_size=1,
                           hvd_id=0)
    input_coors, input_features, input_num_list, input_bbox = [], [], [], []
    for i in tqdm(range(epoch)):
        coors_d, features_d, num_list_d, bbox_d = next(WaymoDataset.train_generator())
        input_coors.append(coors_d)
        input_features.append(features_d)
        input_num_list.append(num_list_d)
        input_bbox.append(bbox_d)
    WaymoDataset.stop()

    coors_p = tf.placeholder(dtype=tf.float32, shape=[None, 3])
    features_p = tf.placeholder(dtype=tf.float32, shape=[None, 3])
    num_list_p = tf.placeholder(dtype=tf.int32, shape=[None])
    bbox_p = tf.placeholder(dtype=tf.float32, shape=[None, 256, 9])
    coors, features, num_list, bbox = coors_p, features_p, num_list_p, bbox_p

    # coors, num_list = grid_sampling(input_coors=coors,
    #                                 input_num_list=num_list,
    #                                 resolution=0.2)

    center_coors, center_num_list, _ = grid_sampling(input_coors=coors,
                                                     input_num_list=num_list,
                                                     resolution=0.4,
                                                     dimension=config.dimension_training,
                                                     offset=config.offset_training)

    roi_attrs, roi_conf, _ = get_roi_bbox(center_coors, bbox, center_num_list, anchor_size)
    roi_conf = tf.cast(roi_conf, dtype=tf.float32)
    roi_attrs, roi_num_list, _ = roi_filter(input_roi_attrs=roi_attrs,
                                            input_roi_conf=roi_conf,
                                            input_roi_ious=roi_conf,
                                            input_num_list=center_num_list,
                                            conf_thres=0.9,
                                            iou_thres=0,
                                            max_length=512,
                                            with_negative=False)

    rois_attrs_noise = tf.random.uniform(shape=[tf.shape(roi_attrs)[0], 7],
                                         minval=-0.1,
                                         maxval=0.1,
                                         dtype=tf.dtypes.float32,
                                         seed=None,
                                         name=None)
    # roi_attrs += rois_attrs_noise


    roi_voxels = la_roi_pooling_fast(input_coors=coors,
                                         input_features=features,
                                         roi_attrs=roi_attrs,
                                         input_num_list=num_list,
                                         roi_num_list=roi_num_list,
                                         dimension=config.dimension_training,
                                         offset=config.offset_training,
                                         grid_buffer_resolution=0.8,
                                         grid_buffer_size=8,
                                         voxel_size=5,
                                         padding_value=0.,
                                         pooling_size=2)
    voxels = voxel2col(roi_voxels, 3)


    init_op = tf.initialize_all_variables()
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = False
    config.allow_soft_placement = False
    config.log_device_placement = False
    config.gpu_options.visible_device_list = '0'
    with tf.Session(config=config) as sess:
        sess.run(init_op)
        for n in tqdm(range(epoch)):

            output_roi_voxels, output_voxels, output_roi_attrs, output_num_list = sess.run([roi_voxels, voxels, roi_attrs, roi_num_list],
            # output_attrs = sess.run(roi_attrs,
                                    feed_dict={coors_p: input_coors[n],
                                               features_p: get_rgbs_from_coors(input_coors[n], repeat=20),
                                               num_list_p: input_num_list[n],
                                               bbox_p: input_bbox[n]})

            output_voxels = np.reshape(output_voxels, newshape=[output_voxels.shape[0], output_voxels.shape[1], 27, -1])
            output = np.zeros([output_voxels.shape[0], output_voxels.shape[1], output_voxels.shape[3]])
            for i in tqdm(range(output_voxels.shape[0] * output_voxels.shape[1])):
                input_num = i // output_voxels.shape[1]
                output_num = i % output_voxels.shape[1]
                non_void_idx = np.sum(output_voxels[input_num, output_num, :, :], axis=-1) > 0
                for j in range(3):
                    output[input_num, output_num, j] = np.mean(output_voxels[input_num, output_num, non_void_idx, j])

    id = 2

    output_voxels = fetch_instance(output, output_num_list, id=id)
    output_roi_attrs = fetch_instance(output_roi_attrs, output_num_list, id=id)
    plot_points_from_roi_voxels(voxels=output_voxels,
                                roi_attrs=output_roi_attrs,
                                kernel_size=3,
                                name='voxel2col')

    output_roi_voxels = fetch_instance(output_roi_voxels, output_num_list, id=id)
    plot_points_from_roi_voxels(voxels=output_roi_voxels,
                                roi_attrs=output_roi_attrs,
                                kernel_size=5,
                                name='voxel2col_roi')

    output_coors = fetch_instance(input_coors[n], input_num_list[n], id=id)
    output_features = fetch_instance(get_rgbs_from_coors(input_coors[n], repeat=20), input_num_list[n], id=id)
    plot_points(output_coors, rgb=output_features, name='voxel2col_input')

