import os
import argparse
import yaml
import json
from utils.dataset import register_box_dataset
from utils.evaluator import Evaluator, create_dirs
from locator_predictor import LocatorPredictor


class EndToEndEvaluator(Evaluator):

    ERROR_RESULT_CODES = ['ERROR', 'TIMEOUT', 'ONLY_OCCLUDED', 'EMPTY_BBOX']

    def __init__(self, dataset_dir, results_dir, pholocator_path, phoxicontrol_path, plcf_path, plcf_settings_json):
        super().__init__(
            results_dir,
            os.path.basename(dataset_dir),
            "locator_predictor",
            step="final",  # "final" or string consisting of 7 decimals
            config=None,  # config
            weights=None,  # weights
            score_threshold=0.5,
            iou_type="segm",  # 'bbox' or 'segm'
            iou_threshold=0.8,
            low_bad_pick_iou=0.1,
            high_bad_pick_iou=0.8)

        self.pholocator_path = pholocator_path
        self.phoxicontrol_path = phoxicontrol_path
        self.plcf_path = plcf_path
        self.plcf_settings_json = plcf_settings_json
        self.error_dt_json = os.path.join(self.step_dir, "error_results.json")

    def gts_by_filename(self):
        result = {}
        category_id = self.cocogt.getCatIds()[0]
        for img_dict in self.cocogt.loadImgs(self.cocogt.getImgIds()):
            img_dict["category_id"] = category_id
            result[os.path.splitext(img_dict['file_name'])[0]] = img_dict
        return result

    @staticmethod
    def load_yaml(yaml_path):
        with open(yaml_path) as f:
            data = yaml.load(f, Loader=yaml.FullLoader)
            return data

    @staticmethod
    def get_segmentation(yaml_data, width, height):
        segmentations = []
        bboxes = []
        if "poses" in yaml_data.keys():
            for pose in yaml_data["poses"]:
                segmentation = []
                bbox = None
                for box_corner in pose["box_corners"]:
                    x = box_corner['x'] * int(width)
                    y = box_corner['y'] * int(height)
                    segmentation.append(x)
                    segmentation.append(y)
                    if bbox is None:
                        bbox = [x, y, x, y]
                    bbox = [min(x, bbox[0]), min(y, bbox[1]), max(x, bbox[2]), max(y, bbox[3])]

                bbox = [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]]
                segmentations.append(segmentation)
                bboxes.append(bbox)

        return bboxes, segmentations

    def create_dt_json(self, yaml_dir):
        gts_by_filename = self.gts_by_filename()
        yaml_files = sorted([f for f in os.listdir(yaml_dir) if os.path.splitext(f)[1] == ".yaml"])

        results_json_data = []
        error_results_json_data = []

        for i, yaml_file in enumerate(yaml_files):
            filename = os.path.splitext(yaml_file)[0]
            yaml_data = self.load_yaml(os.path.join(yaml_dir, yaml_file))
            result_code = yaml_data["result_code"]

            gt = gts_by_filename.get(filename)
            if result_code == 'SUCCESS':
                if gt is None:
                    print("Corresponding ground truth is missing for {} ".format(yaml_file))
                    continue
                bboxes, segmentations = self.get_segmentation(yaml_data, gt['width'], gt['height'])
                for bbox, segmentation in zip(bboxes, segmentations):
                    results_json_data.append({
                        "image_id": gt["id"],
                        "category_id": gt["category_id"],
                        "score": 1.0,
                        "bbox": bbox,
                        "segmentation": [segmentation],
                        "metadata": {}
                    })
            elif result_code in self.ERROR_RESULT_CODES:
                error_results_json_data.append({
                    "error_result_code": result_code,
                    "file_name": "{}.png".format(filename)
                })

        with open(self.dt_json, 'w') as f:
            json.dump(results_json_data, f)

        with open(self.error_dt_json, 'w') as f:
            json.dump(error_results_json_data, f)

    def predict_dataset(self, ):
        predictor = LocatorPredictor(self.metadata.image_root, self.result_dir, self.pholocator_path, self.phoxicontrol_path, self.plcf_path, self.plcf_settings_json)
        #predictor.detect()

        #TODO:: prerobit
        self.create_dt_json(predictor.result_yaml_dir)

    def create_dirs(self):
        super().create_dirs()
        create_dirs([os.path.join(self.export_dir, code.lower()) for code in self.ERROR_RESULT_CODES])

    def debug(self):
        super().debug()
        with open(self.error_dt_json, 'r') as f:
            error_dt = json.load(f)
            for el in error_dt:
                filename = el["file_name"]
                code = el["error_result_code"]
                os.symlink(self.rel_img_source_path(filename), os.path.join(self.export_dir, code.lower(), filename))


if __name__ == "__main__":
    USER = os.environ['USER']
    parser = argparse.ArgumentParser()

    # sat args
    parser.add_argument("--dataset_dir", required=True,
                        help="path to folder with praws, annotation and textures")

    parser.add_argument("--plcf_path", required=True,
                        help="path to folder with plcfs")

    parser.add_argument("--plcf_settings_json", required=True,
                        help="path to json with plcf setting, plcf's will be updated for this config")

    parser.add_argument("--results_dir", default='results/',
                        help="results json and subfolders with plotted results")

    parser.add_argument("--pholocator_path", default='pholocator',
                        help="name of pholocator proces")

    parser.add_argument("--phoxicontrol_path", default='PhoXiControl',
                        help="name of phoxicontrol proces")

    args = parser.parse_args()

    register_box_dataset(args.dataset_dir)
    evaluator = EndToEndEvaluator(args.dataset_dir, args.results_dir, args.pholocator_path, args.phoxicontrol_path, args.plcf_path, args.plcf_settings_json)
    evaluator.predict_dataset()
    evaluator.visualize()
    evaluator.debug()
    evaluator.report([0.75, 0.8, 0.85, 0.9])