#!/usr/bin/env python
from __future__ import print_function
import numpy as np
import sys
import os
import cv2

import ensure_segmappy_is_installed
from segmappy import Config
from segmappy import Dataset
from segmappy import Generator
from segmappy.tools.classifiertools import get_default_dataset, get_default_preprocessor
from segmappy.tools.roccurve import get_roc_pairs, get_roc_curve
from segmappy.models.model_groups_tf import init_model
from segmappy.tools.heatmap import plot_heatmap_img

# read config file
configfile = "default_training.ini"

config = Config(configfile)

# add command line arguments to config
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--log")
parser.add_argument("--heatmap_dir", default="/tmp/segmap/heatmaps/")
args = parser.parse_args()
config.log_name = args.log
config.heatmap_dir = args.heatmap_dir

if not os.path.exists(config.heatmap_dir):
    os.makedirs(config.heatmap_dir)

ranks_data = np.loadtxt("/mnt/data/datasets/JW/MulRan/segmap/training_datasets/DCC03/ranks.log", dtype=np.int)
nbins = 10
hist_sum = np.zeros(nbins, dtype=np.float)
hist_sum_sa = np.zeros(nbins, dtype=np.float)
hist_n = np.zeros(nbins, dtype=np.int)
cid_to_rank = {}
cid_to_bin = {}
for i, row in enumerate(ranks_data):
    cid_to_rank[(row[0], row[1])] = row[2]
    cid_to_bin[(row[0], row[1])] = int(float(i) / ranks_data.shape[0] * nbins)


# load dataset and preprocessor
dataset = Dataset(
    folder=config.cnn_test_folder,
    base_dir=config.base_dir,
    keep_match_thresh=0.7,
    # use_merges=False,
    # use_matches=False,
    # normalize_classes=False,
    # remove_poorly_visible=False,
    largest_vis_view=True
)

segments, positions, classes, n_classes, _, _, _, int_paths, mask_paths, range_paths, _ = dataset.load()
duplicate_classes = dataset.duplicate_classes
duplicate_ids = dataset.duplicate_ids

preprocessor = get_default_preprocessor(config)
preprocessor.init_segments(segments,
                           classes,
                           positions=positions,
                           int_paths=int_paths,
                           mask_paths=mask_paths,
                           range_paths=range_paths)

# calculate all feature sets
feature_sets = []

gen_test = Generator(
    preprocessor,
    np.arange(len(segments)),
    np.max(duplicate_classes) + 1,
    train=False,
    batch_size=config.batch_size,
    shuffle=False,
)
print("Testing with %d segments and %d classes" % (gen_test.n_segments, gen_test.n_classes))

import tensorflow as tf

tf.reset_default_graph()

# restore variable names from previous session
saver = tf.train.import_meta_graph(config.cnn_model_folder_w_vis_views + "/model.ckpt.meta")

# get key tensorflow variables
cnn_graph = tf.get_default_graph()

cnn_input = cnn_graph.get_tensor_by_name("InputScope/input:0")
cnn_input_vis = cnn_graph.get_tensor_by_name("InputScope/input_vis:0")
y_true = cnn_graph.get_tensor_by_name("y_true:0")
training = cnn_graph.get_tensor_by_name("training:0")
scales = cnn_graph.get_tensor_by_name("scales:0")

loss = cnn_graph.get_tensor_by_name("loss:0")
loss_c = cnn_graph.get_tensor_by_name("loss_c:0")
loss_r = cnn_graph.get_tensor_by_name("loss_r:0")

accuracy = cnn_graph.get_tensor_by_name("accuracy:0")
# y_prob = cnn_graph.get_tensor_by_name("y_prob:0")
descriptor = cnn_graph.get_tensor_by_name("OutputScope/descriptor_read:0")
roc_auc = cnn_graph.get_tensor_by_name("roc_auc:0")

#visualization
conv = cnn_graph.get_tensor_by_name("conv3_out:0")
# conv_vis_grad = cnn_graph.get_tensor_by_name("conv5_vis_grad_out:0")
img_heatmap = cnn_graph.get_tensor_by_name("img_heatmap:0")
score = cnn_graph.get_tensor_by_name("score:0")

conv_vis = cnn_graph.get_tensor_by_name("conv5_vis_out:0")
img_heatmap_vis = cnn_graph.get_tensor_by_name("img_heatmap_vis:0")
score_vis = cnn_graph.get_tensor_by_name("score_vis:0")

global_step = cnn_graph.get_tensor_by_name("global_step:0")
update_step = cnn_graph.get_tensor_by_name("update_step:0")
train_op = cnn_graph.get_operation_by_name("train_op")

summary_batch = tf.summary.merge_all("summary_batch")
summary_epoch = tf.summary.merge_all("summary_epoch")
summary_heatmap = tf.summary.merge_all("summary_heatmap")
summary_heatmap_vis = tf.summary.merge_all("summary_heatmap_vis")

# os.environ['CUDA_VISIBLE_DEVICES'] = "1"

config_tf = tf.ConfigProto()
config_tf.gpu_options.allow_growth = True

with tf.Session(config=config_tf) as sess:
    # tensorboard statistics
    if config.log_name:
        test_writer = tf.summary.FileWriter(
            os.path.join(config.log_path, config.log_name, "test")
        )

    # initialize all tf variables
    tf.global_variables_initializer().run()

    saver.restore(sess, config.cnn_model_folder_w_vis_views + "/model.ckpt")

    segment_idx = 0
    # sequence of test batches
    batches = [0] * gen_test.n_batches
    for epoch in range(1):
        train_loss = 0
        train_loss_c = 0
        train_loss_r = 0
        train_accuracy = 0
        train_step = 0

        test_loss = 0
        test_loss_c = 0
        test_loss_r = 0
        test_accuracy = 0
        test_step = 0

        np.random.shuffle(batches)

        console_output_size = 0
        for step, train in enumerate(batches):
            batch_segments, batch_classes, batch_vis_views = gen_test.next()

            # calculate test loss and accuracy
            [batch_conv_vis, batch_descriptor] = sess.run(
                [conv_vis, descriptor],
                feed_dict={
                    cnn_input: batch_segments,
                    cnn_input_vis: batch_vis_views,
                    scales: preprocessor.last_scales,
                },
            )

            # if config.log_name:
            #     test_writer.add_summary(summary, sess.run(global_step))

            if True:
                # img_heatmap_vis_np, val_heatmap_np, scores_vis = plot_heatmap_img(batch_segments,
                #                                                                   batch_vis_views,
                #                                                                   preprocessor.last_scales,
                #                                                                   batch_conv_vis,
                #                                                                   batch_descriptor,
                #                                                                   sess,
                #                                                                   descriptor,
                #                                                                   cnn_input,
                #                                                                   cnn_input_vis,
                #                                                                   scales)
                img_heatmap_vis_np = np.zeros(batch_vis_views.shape, dtype=np.uint8)
                val_heatmap_np = np.zeros(batch_vis_views.shape[:-1], dtype=np.float)
                for i in range(batch_segments.shape[0]):
                    if segment_idx < len(duplicate_classes):
                        cur_id = duplicate_classes[segment_idx]
                        cur_did = duplicate_ids[segment_idx]
                        seg_dir = "%06d" % cur_id
                        seg_filename = "%03d" % cur_did

                        if not os.path.exists(os.path.join(config.heatmap_dir, seg_dir)):
                            os.makedirs(os.path.join(config.heatmap_dir, seg_dir))

                        # cv2.imwrite(os.path.join(config.heatmap_dir, seg_dir, seg_filename + "_heat_img.jpg"),
                        #             cv2.cvtColor(img_heatmap_vis_np[i], cv2.COLOR_BGR2RGB))
                        # cv2.imwrite(os.path.join(config.heatmap_dir, seg_dir, seg_filename + "_heat_val.png"),
                        #             (val_heatmap_np[i]*65535).astype(np.uint16))

                        img_heatmap_vis_np[i] = cv2.cvtColor(
                                cv2.imread(os.path.join(config.heatmap_dir, seg_dir, seg_filename + "_heat_img.jpg")),
                                                             cv2.COLOR_BGR2RGB)
                        val_heatmap_np[i] = cv2.imread(os.path.join(config.heatmap_dir, seg_dir, seg_filename + "_heat_val.png"),
                                                                cv2.IMREAD_ANYDEPTH).astype(np.float) / 65535.0

                        mask = batch_vis_views[i, :, :, 1]
                        # mask and its vicinity
                        mask = cv2.dilate(mask, np.ones((7, 7), np.uint8))
                        mask_area = np.sum(mask > 0.5)
                        val_mask = np.sum(val_heatmap_np[i][mask > 0.5])
                        val_all = np.sum(val_heatmap_np[i])
                        if (cur_id, cur_did) in cid_to_bin:
                            bin_idx = cid_to_bin[(cur_id, cur_did)]
                            hist_sum[bin_idx] += val_mask / val_all
                            hist_sum_sa[bin_idx] += (val_mask / mask_area) / \
                                                 (val_all / (mask.shape[0]*mask.shape[1] - mask_area))
                            hist_n[bin_idx] += 1

                        segment_idx += 1
                    else:
                        print(segment_idx)
                        print(len(duplicate_classes))
                        print(batch_segments.shape)
                        print(img_heatmap_vis_np.shape)

                # img_heatmap_np, scores = plot_heatmap_vox(batch_segments,
                #                                   batch_vis_views,
                #                                   batch_classes,
                #                                   preprocessor.last_scales,
                #                                   batch_conv,
                #                                   batch_descriptor,
                #                                   sess,
                #                                   descriptor,
                #                                   cnn_input,
                #                                   cnn_input_vis,
                #                                   y_true,
                #                                   scales,
                #                                   480, 640)

                # summary = sess.run(summary_heatmap, feed_dict={img_heatmap_vis: img_heatmap_vis_np,
                #                                                img_heatmap: img_heatmap_np,
                #                                                score: scores,
                #                                                score_vis: scores_vis})
                # summary = sess.run(summary_heatmap_vis, feed_dict={img_heatmap_vis: img_heatmap_vis_np,
                #                                                    score_vis: scores_vis})
                # test_writer.add_summary(summary, sess.run(global_step))
                # test_writer.flush()

            # test_loss += batch_loss
            # test_loss_c += batch_loss_c
            # test_loss_r += batch_loss_r
            # test_accuracy += batch_accuracy
            test_step += 1

            # print results
            sys.stdout.write("\b" * console_output_size)

            console_output = "epoch %2d, step %05d " % (epoch, step)

            if test_step:
                console_output += "v_loss %.4f v_acc %.2f v_c %.4f v_r %.4f" % (
                    test_loss / test_step,
                    test_accuracy / test_step * 100,
                    test_loss_c / test_step,
                    test_loss_r / test_step,
                )

            console_output_size = len(console_output)

            sys.stdout.write(console_output)
            sys.stdout.flush()

        # flush tensorboard log
        if config.log_name:
            test_writer.flush()

        # # save epoch model
        # if not config.keep_best or test_accuracy > best_accuracy:
        #     if config.checkpoints > 1:
        #         model_name = "model-%d.ckpt" % sess.run(global_step)
        #     else:
        #         model_name = "model.ckpt"
        #
        #     if config.vis_views:
        #         model_folder = config.cnn_model_folder_w_vis_views
        #     else:
        #         model_folder = config.cnn_model_folder_wo_vis_views
        #     saver.save(sess, os.path.join(model_folder, model_name))
        #     tf.train.write_graph(
        #         sess.graph.as_graph_def(), model_folder, "graph.pb"
        #     )

        hist_n = np.maximum(hist_n, 1)

        print(hist_sum / hist_n)
        print(hist_sum_sa / hist_n)
        print(hist_n)
