import numpy as np
import os
import fiftyone as fo


def export_yolo_data(samples,
                     export_dir,
                     classes,
                     label_field="ground_truth",
                     split=None):

    if type(split) == list:
        splits = split
        for split in splits:
            export_yolo_data(samples, export_dir, classes, label_field, split)
    else:
        if split is None:
            split_view = samples
            split = "val"
        else:
            split_view = samples.match_tags(split)

        split_view.export(export_dir=export_dir,
                          dataset_type=fo.types.YOLOv5Dataset,
                          label_field=label_field,
                          classes=classes,
                          split=split)


def read_yolo_detections_file(filepath):
    detections = []
    if not os.path.exists(filepath):
        return np.array([])

    with open(filepath) as f:
        lines = [line.rstrip('\n').split(' ') for line in f]

    for line in lines:
        detection = [float(l) for l in line]
        detections.append(detection)
    return np.array(detections)


def _uncenter_boxes(boxes):
    '''convert from center coords to corner coords'''
    boxes[:, 0] -= boxes[:, 2] / 2.
    boxes[:, 1] -= boxes[:, 3] / 2.


def _get_class_labels(predicted_classes, class_list):
    labels = predicted_classes.detach().cpu().int()
    labels = [class_list[l] for l in labels]
    return labels


def convert_yolo_segmentations_to_fiftyone(yolo_segmentations, class_list):

    detections = []
    boxes = yolo_segmentations.boxes.xywhn
    if not boxes.shape or yolo_segmentations.masks is None:
        return fo.Detections(detections=detections)

    confs = yolo_segmentations.boxes.conf.detach().cpu().numpy()
    _uncenter_boxes(boxes)
    masks = yolo_segmentations.masks.masks
    labels = _get_class_labels(yolo_segmentations.boxes.cls, class_list)

    for label, box, conf, mask in zip(labels, boxes, confs, masks):
        ## convert to absolute indices to index mask
        w, h = mask.shape
        box = box.detach().cpu()
        mask = mask.detach().cpu().numpy()
        tmp = np.copy(box)
        tmp[2] += tmp[0]
        tmp[3] += tmp[1]
        tmp[0] *= h
        tmp[2] *= h
        tmp[1] *= w
        tmp[3] *= w
        tmp = [int(b) for b in tmp]
        y0, x0, y1, x1 = tmp
        sub_mask = mask[x0:x1, y0:y1]

        detections.append(
            fo.Detection(label=label,
                         bounding_box=list(box),
                         mask=sub_mask,
                         confidence=conf))

    return fo.Detections(detections=detections)


def convert_yolo_detections_to_fiftyone(yolo_detections, class_list):

    detections = []
    if yolo_detections.size == 0:
        return fo.Detections(detections=detections)

    boxes = yolo_detections[:, 1:-1]
    _uncenter_boxes(boxes)

    confs = yolo_detections[:, -1]
    labels = _get_class_labels(yolo_detections[:, 0], class_list)

    for label, conf, box in zip(labels, confs, boxes):
        detections.append(
            fo.Detection(label=label,
                         bounding_box=box.tolist(),
                         confidence=conf))

    return fo.Detections(detections=detections)


def get_prediction_filepath(filepath, run_number=1):
    run_num_string = ""
    if run_number != 1:
        run_num_string = str(run_number)
    filename = filepath.split("/")[-1].split(".")[0]
    return f"runs/detect/predict{run_num_string}/labels/{filename}.txt"


def add_yolo_detections(samples, prediction_field, prediction_filepath,
                        class_list):

    prediction_filepaths = samples.values(prediction_filepath)
    yolo_detections = [
        read_yolo_detections_file(pf) for pf in prediction_filepaths
    ]
    detections = [
        convert_yolo_segmentations_to_fiftyone(yd, class_list)
        for yd in yolo_detections
    ]
    samples.set_values(prediction_field, detections)

    
def order_detections(detections):
    """Takes Detection objects with bounding boxes in 
    [<top-left-x>, <top-left-y>, <width>, <height>]
    format and orders them from top to bottom.
    
    :returns: ordered list of Detection objects
    """
    result = detections.copy()
    result.sort(key=sort_t_b)
    return result
    

def sort_t_b(detection):
    """Passed to sort() function as key to order detections top
    to bottom.
    """
    return detection['bounding_box'][1]