# This file contains code to plot F-ROC and O-ROC file for the baseline
# It can be used for either all of them or just one of them (detection and recognition)

import logging
import yamlparser
import argparse
import os
import pickle
import numpy as np
from ..dataset import read_detections,read_ground_truth,read_recognitions
from ..evaluate import assign_detections,compute_DR_FDPI,plot_froc_curve,compute_TPIR_FPIPI,plot_oroc_curve

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("FaceRec.UCCS")

def read_config_file():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--tasks", "-t",
        nargs = "+",
        choices = ["detection", "recognition"],
        default = ["detection", "recognition"],
        help = "Select the tasks that should be performed in this evaluation"
    )

    parent_direct = os.path.dirname(os.path.dirname(__file__))
    cfg = yamlparser.config_parser(parser=parser,default_config_files=[os.path.join(parent_direct, "configs/baseline_config.yaml")])

    if 'detection' in cfg.tasks:
        if cfg.eval.detection.files is None:
            raise ValueError("For the detection task, --eval.detection.files are required.")

        if cfg.eval.detection.labels is None:
            cfg.eval.detection.labels = cfg.eval.detection.files

        if len(cfg.eval.detection.labels) != len(cfg.eval.detection.files):
            raise ValueError("The number of labels (%d) and results (%d) differ" % (len(cfg.eval.detection.labels), len(cfg.eval.detection.files)))

    if 'recognition' in cfg.tasks:
        if cfg.eval.recognition.files is None:
            raise ValueError("For the identification task, --eval.recognition.files are required.")

        if cfg.eval.recognition.labels is None:
            cfg.eval.recognition.labels = cfg.eval.recognition.files

        if len(cfg.eval.recognition.labels) != len(cfg.eval.recognition.files):
            raise ValueError("The number of labels (%d) and results (%d) differ" % (len(cfg.eval.recognition.labels), len(cfg.eval.recognition.labels)))

    return cfg

def main():

    # get config arguments
    cfg = read_config_file()

    # load the evaluation protocol
    logger.info("Loading UCCS %s evaluation protocol",cfg.which_set)

    #read detections based on the task (detection/identification)
    if "detection" in cfg.tasks or "recognition" in cfg.tasks:
        # read the ground truth bounding boxes of the validation set
        logger.info("Reading UCCS %s ground-truth from the protocol",cfg.which_set)
        ground_truth = read_ground_truth(cfg.data_directory,cfg.which_set)

        face_numbers = sum([len(face_ids) for face_ids,_,_ in ground_truth.values()])
        image_numbers = len(ground_truth)

        # exclude gallery faces from the result
        if cfg.eval.exclude_gallery is not None:
            
            with open(cfg.eval.exclude_gallery, 'r') as file:
                exclude = [int(line.strip()) for line in file]

            # update the number of faces in the set
            face_numbers -= len(exclude)

    # plot f-roc for the detection results if it is given
    if  "detection" in cfg.tasks:

        detection_results = []

        for idx,detection_file in enumerate(cfg.eval.detection.files):

            logger.info("Reading detections from %s (%s)", detection_file, cfg.eval.detection.labels[idx])
            detections = read_detections(detection_file)

            matched_detections = assign_detections(ground_truth,detections,cfg.eval.iou,exclude)

            logger.info("Computing DR and FDPI for %s (%s)", detection_file, cfg.eval.detection.labels[idx])
            DR,FDPI = compute_DR_FDPI(matched_detections,face_numbers,image_numbers,cfg.eval.detection.plot_numbers)

            detection_results.append((DR,FDPI))

        # plotting
        froc_path = cfg.eval.detection.froc
        logger.info("Plotting F-ROC curve(s) to file '%s'", froc_path)
        plot_froc_curve(detection_results,cfg.eval.detection.labels,froc_path,
                                 face_numbers,cfg.eval.linear,cfg.eval.detection.plot_numbers)

    # plot o-roc for the identification results if it is given
    if "recognition" in cfg.tasks:

        known_numbers = sum([ sum(np.array(subject_ids) > 0) for _,subject_ids,_ in ground_truth.values()])
        if exclude:
            known_numbers -= len(exclude)

        recognition_results = []

        for idx, score_file in enumerate(cfg.eval.recognition.files):
            logger.info("Reading scores from %s (%s)", score_file, cfg.eval.recognition.labels[idx])
            all_scores = read_recognitions(score_file)

            matched_detections = assign_detections(ground_truth,all_scores,cfg.eval.iou,exclude)

            logger.info("Computing TPIR and FPIPI for %s (%s)", score_file, cfg.eval.recognition.labels[idx])
            TPIR,FPIPI = compute_TPIR_FPIPI(all_scores,matched_detections,known_numbers,image_numbers,
                                                    cfg.eval.recognition.rank,cfg.eval.recognition.plot_numbers)

            recognition_results.append((TPIR,FPIPI))

        # plotting
        oroc_path = cfg.eval.recognition.oroc
        logger.info("Plotting O-ROC curve(s) to file '%s'", oroc_path)
        plot_oroc_curve(recognition_results,cfg.eval.recognition.labels,cfg.eval.recognition.rank,oroc_path,
                                 known_numbers,linear=cfg.eval.linear,
                                 plot_recognition_numbers=cfg.eval.recognition.plot_numbers)

if __name__ == "__main__":
    main()
