import sys
sys.path.append("/workspaces/tensorflow-gpu/cocoapi/PythonAPI/pycocotools")
import os
import cv2
from detectron2.data import DatasetCatalog
import detectron2.data as ddata
from detectron2.engine import DefaultPredictor
from detectron2.evaluation import COCOEvaluator
from detectron2.utils.visualizer import ColorMode
from detectron2.utils.visualizer import Visualizer
import random
import matplotlib.pyplot as plt
import time
from from_root import from_root
from rcnn_model.utils.floorplan_vectorizer_utils import check_image_size_thresh


results_directory = str(from_root("results"))+"/"
max_image_size = 700*500

def main(cfg,results_filename = "eval_results.txt"):
    #update config file
    cfg.DATALOADER.NUM_WORKERS = 1
    cfg.SOLVER.IMS_PER_BATCH = 1                                           
    cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")

    #run evaluation
    results = standard_evaluation(cfg)

    #save results
    file = open(results_directory+results_filename,"w")
    file.write(str(results))
    file.close()


### Evaluation ###

def standard_evaluation(cfg):
    #load predictor
    predictor = DefaultPredictor(cfg)
    test_data_loader = ddata.build_detection_test_loader(cfg, "inodata_val")

    #save some validation images
    save_validation_images(predictor)

    #create evaluator
    evaluator = COCOEvaluator("inodata_val",tasks={"segm","bbox"},output_dir="./eval_output",distributed=False,max_dets_per_image=50,allow_cached_coco=False)
    print("EVALUATING")
    evaluator.reset()

    #load results into evaluator
    for inputs, outputs in block_prediction(test_data_loader, predictor):
        evaluator.process(inputs,outputs)
        del inputs
        del outputs
        print("|",end="")
        time.sleep(.5)
    print("")

    #run evaluator
    results = evaluator.evaluate()
    print(results)
    print("EVALUATED")
    return results


def block_prediction(loader, predictor):
    for data in loader:
        if(check_image_size_thresh(data[0]["file_name"],max_image_size)):
            image = cv2.imread(data[0]["file_name"])
            result = predictor(image)
            yield data, [result]
            del image
            del result


### Validation Images ###

def save_validation_images(predictor):
    val_img_id = 1
    for d in random.sample(DatasetCatalog.get("inodata_val"), 16):
        save_image(d,predictor,val_img_id)
        val_img_id += 1


def save_image(d, predictor, val_img_id):
    try:
        if(check_image_size_thresh(d["file_name"],max_image_size)):
            #set load image
            val_img_dest_path = "models/rcnn/validation_images/RCNN_val_image_"+str(val_img_id)+".png"
            im = cv2.imread(d["file_name"])
            outputs = predictor(im)

            #save image
            v = Visualizer(im[:,:,::-1],scale=0.5,instance_mode=ColorMode.IMAGE_BW)
            out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
            plt.imshow(out.get_image()[:,:,::-1])
            plt.axis('off')
            plt.savefig(val_img_dest_path,bbox_inches='tight',pad_inches=0)
            print("Saved validation image to "+val_img_dest_path)
            plt.clf()
    except:
        print("ERROR SAVING IMAGE")