'''
Class for detecting flowers and flower pots in 2D image
'''

import rospy

import torch
import torch.backends.cudnn as cudnn
import numpy as np
import time
from pathlib import Path

# yolov5 modules
from models.experimental import attempt_load
from utils.augmentations import letterbox
from utils.general import check_img_size, non_max_suppression, scale_coords, set_logging
from utils.plots import colors, Annotator
from utils.torch_utils import select_device, time_sync

FILE = Path(__file__).resolve()
ROOT = FILE.parents[2]

class YoloDetector:

    def __init__(self, device_num):

		### launch parameters
        try:
            YOLO_MODEL_PATH = str(rospy.get_param('yolo_model_path'))
            IMG_SIZE = int(rospy.get_param('img_size'))
            MIN_CONF = float(rospy.get_param('min_conf'))
        except Exception:
            print("Can't load Yolo parameters from ROS parameter server, using defaults.")
            YOLO_MODEL_PATH = 'data/best_real.pt'
            IMG_SIZE = 416
            MIN_CONF = 0.40

        ### init atributes
        self.DO_DEBUGING = False

        self.xyxy = None
        self.conf = None
        self.cls = None

        ### yolov5 initialization - TODO vymazat nepotrebne parametry

        weights = ROOT / YOLO_MODEL_PATH    # model.pt path(s)
        self.min_conf = MIN_CONF            # minimum confidence to consider a detection
        self.imgsz = IMG_SIZE               # inference size (pixels)
        self.conf_thres = MIN_CONF-0.01     # confidence threshold for NMS
        self.iou_thres = 0.45               # NMS IOU threshold # 0.45
        self.max_det = 100                  # maximum detections per image
        self.classes = None                 # filter by class: --class 0, or --class 0 2 3
        self.agnostic_nms = False # False!  # class-agnostic NMS (opposite of class-aware)
        self.line_width = 1                 # bounding box thickness (pixels)
        self.hide_labels = False            # hide labels 
        self.hide_conf = False              # hide confidences
        self.half = False # False!          # use FP16 half-precision (o 20 % pomalejsi)
        self.stride = 32
        # self.device_num = ''              # cuda device, i.e. 0 or 0,1,2,3 or cpu

        # Initialize
        set_logging()
        self.device = select_device(device_num)
        self.half &= self.device.type != 'cpu'  # half precision only supported on CUDA

        # Load model
        self.model = attempt_load(weights, map_location=self.device)  # load FP32 model
        stride = int(self.model.stride.max())  # model stride
        self.imgsz = check_img_size(self.imgsz, s=stride)  # check image size
        self.names = self.model.module.names if hasattr(self.model, 'module') \
            else self.model.names  # get class names
        if self.half:
            self.model.half()  # to FP16

        # Dataloader
        cudnn.benchmark = True  # speed up constant img size inference (o par %)

        # Run inference once - first detection takes around 15-20 s on Jetson Nano
        if self.device.type != 'cpu':   # self.imgsz nebo 416 x 256 (stejny efekt)
            self.model(torch.zeros(1, 3, self.imgsz, self.imgsz).to( \
                self.device).type_as(next(self.model.parameters())))


    '''
    detect plants and flower pots
        Output:
            flag_anything
    '''
    def perform_detection(self,img0):

        ### INITIALIZE FLAGS
        self.flag_anything = False     # any object
        self.flag_any_flower = False   # any flower
        self.flag_flower = False       # flower within min_conf threshold
        self.flag_pot = False          # pot within min_conf threshold
        self.flag_pair = False          # pot within min_conf threshold

        ### IMAGE PREPARATION

        t0 = time.time()

        # get image
        self.img0 = img0.copy() # image with original resolution

        # Letterbox - pad image with gray to preserve stride-multiple aspect ratio
        img = letterbox(self.img0, self.imgsz, stride=self.stride, auto=True)[0]
        img = img[np.newaxis, :, :, :]    

        # Prepare image for pytorch
        img = np.stack(img, 0)
        img = img[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.half() if self.half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        t1 = time_sync()
        dt1 = np.round(1000*(t1-t0))

        ### INFERENCE AND BASIC PROCESSING

        # Inference
        with torch.no_grad(): # TODO zjisti, zda no_grad na Jetsonu pomaha
            pred = self.model(img)[0]

        t2 = time_sync()
        dt2 = np.round(1000*(t2-t1))

        # Apply NMS
        pred = non_max_suppression(pred, self.conf_thres, self.iou_thres,
            self.classes, self.agnostic_nms, max_det=self.max_det)

        t3 = time_sync()
        dt3 = np.round(1000*(t3-t2))

        if self.DO_DEBUGING:
            rospy.loginfo(f'Preprocessing: {dt1} ms; Detection: {dt2} ms; NMS: {dt3} ms.')

        ### BASIC PROCESSING OF DETECTIONS - returns vectors of BBoxes, confidences, classes
        
        self.det = pred[0]
        if len(self.det):

            # Rescale boxes from img_size to im0 size
            self.det[:,:4] = scale_coords(img.shape[2:], self.det[:,:4], self.img0.shape).round()

            self.flag_anything = True

            self.xyxy = self.det[:,:4].detach().cpu().numpy()
            self.conf = self.det[:,4].detach().cpu().numpy()
            self.cls = self.det[:,5].detach().cpu().numpy().astype('int')


    '''
    Sort detections acording to classes if flower with high enough confidence detected
        Output:
            flag_any_flower
            flag_flower
            flag_pot
    '''
    def process_detections(self):

        flower_idxs = np.nonzero(self.cls == 0)[0]
        pot_idxs = np.nonzero(self.cls == 1)[0]

        if flower_idxs.size: # at least one flower detected

            self.flag_any_flower = True

            self.flowers_conf = self.conf[flower_idxs]
            max_flower_conf = np.amax(self.flowers_conf)
            max_flower_conf_idx = np.argmax(self.flowers_conf) # TODO smaz

            if max_flower_conf >= self.min_conf:

                self.flag_flower = True

                self.flowers_xyxy = self.xyxy[flower_idxs]
                self.flowers_cls = self.cls[flower_idxs] # TODO smaz?

                if pot_idxs.size: # at least one pot detected

                    self.flag_pot = True

                    self.pots_xyxy = self.xyxy[pot_idxs]
                    self.pots_conf = self.conf[pot_idxs]
                    self.pots_cls = self.cls[pot_idxs] # TODO smaz?
                    max_pot_conf = np.amax(self.pots_conf) # TODO smaz
                    max_pot_conf_idx = np.argmax(self.pots_conf) # TODO smaz


    '''
    Form flower-pot pairs from existing detections
        Output:
            np.array[flower_idx, pot_idx] shape: (pair_count,2).
            Indexes pointing to flowers_xyxy ..., pots_xyxy ...
    '''
    def form_pairs(self):

        ### initialization
        # 1 row ~ 1 pair => [flwr_idx, pot_idx]
        pair_idxs = np.zeros((len(self.flowers_conf),2),dtype=np.int)
        pair_count = 0

        # coords of flowers - area of a flower A
        x1f = self.flowers_xyxy[:,0]
        y1f = self.flowers_xyxy[:,1]
        x2f = self.flowers_xyxy[:,2]
        y2f = self.flowers_xyxy[:,3]
        # coords of pots - area of a pot B
        x1p = self.pots_xyxy[:,0]
        y1p = self.pots_xyxy[:,1]
        x2p = self.pots_xyxy[:,2]
        y2p = self.pots_xyxy[:,3]

        # As = (x2f - x1f + 1) * (y2f - y1f + 1) # area of all flowers
        Bs = (x2p - x1p + 1) * (y2p - y1p + 1) # area of all pots

        flower_idxs = np.arange(len(self.flowers_conf)) ## np.argsort(y2f)
        pot_idxs = np.arange(len(self.pots_conf)) ## np.argsort(y2p)

        # for every flower
        for a in flower_idxs:

            ### compute metric as union of the flower with all pots over area of the pots
            # verticies of union
            xx1 = np.maximum(x1f[a], x1p[pot_idxs])
            yy1 = np.maximum(y1f[a], y1p[pot_idxs])
            xx2 = np.minimum(x2f[a], x2p[pot_idxs])
            yy2 = np.minimum(y2f[a], y2p[pot_idxs])
            # sizes of union
            w = np.maximum(0, xx2 - xx1 + 1)
            h = np.maximum(0, yy2 - yy1 + 1)
            # metric = (A union Bi) / Bi
            metric = (w * h) / Bs

            ### choose most suitable pot 
            metric_idxs = np.argsort(metric) # ascending
            metric_max = metric[metric_idxs[-1]]
            if metric_max > 0.5: # TODO maybe different threshold greater than 0
                
                # create new pair
                pair_idxs[pair_count,:] = [ a, pot_idxs[metric_idxs[-1]] ]
                pair_count += 1

                # pop up the pot from available ones
                pot_idxs = np.delete(pot_idxs, metric_idxs[-1])
                Bs = np.delete(Bs, metric_idxs[-1])

                # if no pots left
                if not len(pot_idxs):
                    break
        
        if pair_count > 0:
            self.flag_pair = True 
            self.flower_pot_pairs = pair_idxs[:pair_count,:]
            self.paired_pots_xyxy = self.pots_xyxy[ self.flower_pot_pairs[:,1], : ]
        else:
            self.flower_pot_pairs = np.zeros((0,2),dtype=np.float)


    '''
    Anotate image with unsorted detections
    '''
    def anotate_raw(self):
        
        annotator = Annotator(self.img0, line_width=self.line_width, example=str(self.names))

        for *xyxy, conf, cls in reversed(self.det):
            c = int(cls)  # integer class
            label = None if self.hide_labels else (self.names[c] if \
                self.hide_conf else f'{self.names[c][0]} : {conf:.2f}')
            annotator.box_label(xyxy, label, color=colors(c, True))


    '''
    Anotate image with detected flower-pot pairs
    '''
    def anotate_pairs(self):

        annotator = Annotator(self.img0, line_width=self.line_width, example=str(self.names))
        
        num_pairs = self.flower_pot_pairs.shape[0]

        for i in range(num_pairs-1, -1, -1):
            # pot
            label = (f'p{i}: {self.pots_conf[self.flower_pot_pairs[i,1]]:.2f}')
            annotator.box_label(self.pots_xyxy[self.flower_pot_pairs[i,1]],
                label, color=colors(1, True))
            # flower
            label = (f'f{i}: {self.flowers_conf[self.flower_pot_pairs[i,0]]:.2f}')
            annotator.box_label(self.flowers_xyxy[self.flower_pot_pairs[i,0]],
                label, color=colors(0, True))


    '''
    Anotate image with detected flower-pot pairs and estimated distances of pots
    '''
    def anotate_pairs_distances(self, distances):

        annotator = Annotator(self.img0, line_width=self.line_width, example=str(self.names))
        
        num_pairs = self.flower_pot_pairs.shape[0]

        for i in range(num_pairs-1, -1, -1):
            # pot
            label = (f'p{i}: {self.pots_conf[self.flower_pot_pairs[i,1]]:.2f}; {distances[i][0]:.2f} m')
            annotator.box_label(self.pots_xyxy[self.flower_pot_pairs[i,1]],
                label, color=colors(1, True))
            # flower
            label = (f'f{i}: {self.flowers_conf[self.flower_pot_pairs[i,0]]:.2f}')
            annotator.box_label(self.flowers_xyxy[self.flower_pot_pairs[i,0]],
                label, color=colors(0, True))


