"""Script to spot cat faces in videos and draw bounding boxes around them.
Expects file 'model.best.tar' to exist (generated by train.py).
Writes outputs to outputs/videos/ by default."""
from __future__ import print_function, division
import argparse
import numpy as np
import os
from collections import defaultdict
from scipy import misc
from model import Model2
from common import to_aspect_ratio_add, draw_heatmap, imresize_sidelen
from bbs import RectangleOnImage
import cv2
import torch
from torch.autograd import Variable
from skimage import morphology
#from sklearn.cluster import DBSCAN
import imgaug as ia
import time

torch.backends.cudnn.benchmark = True

WRITE_TO_BASEDIR = "outputs/videos/"
GPU = 0

def main():
    """Find bounding boxes in a video."""

    parser = argparse.ArgumentParser(description="Process a video")
    parser.add_argument("--video", help="Filepath to the video", required=True)
    parser.add_argument("--out_dir", help="Directory name in which to save results")
    parser.add_argument("--start_frame", default=1, type=int, help="Frame number to start at (1 to N).")
    parser.add_argument("--end_frame", help="Frame number to end at (1 to N).")
    parser.add_argument("--conf", default=0.5, type=float, help="Confidence threshold for BBs")
    parser.add_argument("--size", default=400, type=int, help="Input image size when feeding into the model")
    args = parser.parse_args()

    # load trained model
    checkpoint = torch.load("model.best.tar")
    model = Model2()
    model.load_state_dict(checkpoint["state_dict"])
    if GPU >= 0:
        model.cuda(GPU)
    model.eval()
    del checkpoint

    # check if video file exists
    video_fp = args.video
    assert os.path.isfile(video_fp)

    # convert video filename to output directory path
    video_fn = os.path.basename(video_fp)
    if args.out_dir is not None:
        write_to_dir = os.path.join(WRITE_TO_BASEDIR, args.out_dir)
    else:
        write_to_dir = os.path.join(WRITE_TO_BASEDIR, os.path.splitext(video_fn)[0])

    # create output directory if necessary
    if not os.path.exists(write_to_dir):
        os.makedirs(write_to_dir)

    # start reading video
    vidcap = cv2.VideoCapture(video_fp)
    success, img = vidcap.read() # img => float 0-255 BGR
    frame_idx = 0

    # forward video, if --start_frame=<int> was used
    if int(args.start_frame) > 1:
        while frame_idx+1 < int(args.start_frame):
            success, img = vidcap.read() # img => float 0-255 BGR
            frame_idx += 1

    # handle frames of video
    while success:
        # end if args.end_frame=<int> was used and that frame was reached
        if args.end_frame is not None and frame_idx >= args.end_frame:
            break

        # find BBs in frame
        time_start = time.time()
        img_rgb = img[:, :, ::-1]
        time_model = process_frame(frame_idx, img_rgb, model, write_to_dir, args.conf, input_size=args.size)
        #debug_frame(frame_idx, img_rgb, model, args.conf, input_size=args.size)
        time_req = time.time() - time_start

        # output message and forward to next frame
        print("Frame %05d in %03dms (model: %03dms)." % (frame_idx, time_req*1000, time_model*1000))
        success, img = vidcap.read()
        frame_idx += 1

def process_frame(frame_idx, img, model, write_to_dir, conf_threshold, input_size=224):
    """Finds bounding boxes in a video frame, draws these bounding boxes
    and saves the result to HDD.
    """
    # find BBs in frame
    bbs, time_model = find_bbs(img, model, conf_threshold, input_size=input_size)

    # draw BBs
    img_out = np.copy(img)
    for (bb, score) in bbs:
        if score > conf_threshold and bb.width > 2 and bb.height > 2:
            img_out = bb.draw_on_image(img_out, color=[0, 255, 0], thickness=3)

    # save to output directory
    save_to_fp = os.path.join(write_to_dir, "%05d.jpg" % (frame_idx,))
    misc.imsave(save_to_fp, img_out)

    return time_model

def find_bbs(img, model, conf_threshold, input_size):
    """Find bounding boxes in an image."""
    # pad image so that its square
    img_pad, (pad_top, pad_right, pad_bottom, pad_left) = to_aspect_ratio_add(img, 1.0, return_paddings=True)

    # resize padded image to desired input size
    # "linear" interpolation seems to be enough here for 400x400 or larger images
    # change to "area" or "cubic" for marginally better quality
    img_rs = ia.imresize_single_image(img_pad, (input_size, input_size), interpolation="linear")

    # convert to torch-ready input variable
    inputs_np = (np.array([img_rs])/255.0).astype(np.float32).transpose(0, 3, 1, 2)
    inputs = torch.from_numpy(inputs_np)
    inputs = Variable(inputs, volatile=True)
    if GPU >= 0:
        inputs = inputs.cuda(GPU)

    # apply model and measure the model's time
    time_start = time.time()
    outputs_pred = model(inputs)
    time_req = time.time() - time_start

    # process the model's output (i.e. convert heatmaps to BBs)
    result = ModelResult(
        outputs_pred,
        inputs_np,
        img,
        (pad_top, pad_right, pad_bottom, pad_left)
    )
    bbs = result.get_bbs()

    return bbs, time_req

class ModelResult(object):
    """Class the handles the transformation from heatmaps (model output) to
    bounding boxes."""

    def __init__(self, outputs, inputs_np, img, paddings):
        self.inputs = inputs_np
        self.outputs = outputs.cpu().data.numpy()
        assert self.inputs.ndim == 4
        assert self.outputs.ndim == 4
        assert self.inputs.shape[0] == 1
        assert self.outputs.shape[0] == 1
        self.img = img
        self.paddings = paddings
        self.shrink_depth = 1
        self.shrink_threshold = 0.9
        self.heatmap_activation_threshold = 0.25

    def get_bbs(self):
        """Convert model outputs to bounding boxes."""
        outputs_pred = self.outputs

        # generate shape of model input image
        # (=> original image + padding to square it + resize)
        # note: self.inputs has form NCHW
        img_pad_rs_shape = (
            self.inputs.shape[2],
            self.inputs.shape[3],
            3
        )

        # generate shape of original image after padding (no resize)
        img_pad_shape = (
            self.img.shape[0] + self.paddings[0] + self.paddings[2],
            self.img.shape[1] + self.paddings[1] + self.paddings[3],
            3
        )

        # convert heatmaps to rectangles
        # (old code that uses all heatmaps)
        """
        hm_idx_to_rects = []
        for i in range(outputs_pred.shape[1]):
            hms = self._heatmap_to_rects(
                outputs_pred[0, i, ...],
                img_pad_rs_shape
            )

            hms_rev = self._rects_reverse_projection(
                hms, self.img.shape, img_pad_shape,
                self.paddings[0], self.paddings[1],
                self.paddings[2], self.paddings[3]
            )
            hm_idx_to_rects.append(hms_rev)
        bbs = merge_rects_to_bbs(hm_idx_to_rects, conf_threshold)
        """

        # convert only the first heatmap outputs to BBs
        # the other heatmaps (top left corner, top center, ...)
        # are currently ignored
        hm_idx_to_rects = []
        for i in [0]:
            hms = self._heatmap_to_rects(
                outputs_pred[0, i, ...],
                img_pad_rs_shape
            )

            hms_rev = self._rects_reverse_projection(
                hms, self.img.shape, img_pad_shape,
                self.paddings[0], self.paddings[1],
                self.paddings[2], self.paddings[3]
            )
            hm_idx_to_rects.append(hms_rev)
        bbs = hm_idx_to_rects[0]

        return bbs

    def _heatmap_to_rects(self, grid_pred, bb_img):
        """Convert a heatmap to rectangles / bounding box candidates."""
        grid_pred = np.squeeze(grid_pred) # (1, H, W) => (H, W)

        # remove low activations
        grid_thresh = grid_pred >= self.heatmap_activation_threshold

        # find connected components
        grid_labeled, num_labels = morphology.label(
            grid_thresh, background=0, connectivity=1, return_num=True
        )

        # for each connected components,
        # - draw a bounding box around it,
        # - shrink the bounding box to optimal size
        # - estimate a score/confidence value
        bbs = []
        for label in range(1, num_labels+1):
            (yy, xx) = np.nonzero(grid_labeled == label)
            min_y, max_y = np.min(yy), np.max(yy)
            min_x, max_x = np.min(xx), np.max(xx)
            rect = RectangleOnImage(x1=min_x, x2=max_x+1, y1=min_y, y2=max_y+1, shape=grid_labeled)
            activation = self._rect_to_score(rect, grid_pred)
            rect_shrunk, activation_shrunk = self._shrink(grid_pred, rect)
            rect_rs_shrunk = rect_shrunk.on(bb_img)
            bbs.append((rect_rs_shrunk, activation_shrunk))

        return bbs

    def _shrink(self, heatmap, rect):
        """Shrink a rectangle to get rid of some low activations.

        The model often generates areas of high activations, with a few
        pixels of medium activations on the side. When drawing a bounding box
        around these activations, the medium ones can force the bounding box
        to become significantly larger than it should be. This function tries
        to shrink those bounding boxes, while retaining most of the activation.

        This function is implemented in a (slow) recursive way. Using dynamic
        programming would probably be faster.
        """
        assert rect.width >= 1 and rect.height >= 1
        #print("shrink...", rect)
        score_orig = self._rect_to_score(rect, heatmap)
        candidates = self._shrink_candidates(rect, depth=self.shrink_depth)
        candidates_scored = []
        #print("score..")
        for candidate in candidates:
            score = self._rect_to_score(candidate, heatmap)
            score_rel = score / score_orig
            if score_rel >= self.shrink_threshold:
                candidates_scored.append((candidate, score, candidate.area))
        #print("sort ", len(candidates_scored))
        candidates_scored = sorted(candidates_scored, key=lambda t: t[2])
        return (candidates_scored[0][0], candidates_scored[0][1])

    def _shrink_candidates(self, rect, depth):
        """Recursive function called by _shrink() to generate bounding box
        candidates that are smaller than the input bounding box."""
        result = [rect]

        if depth > 0:
            if rect.width > 1:
                rect_left = rect.copy(x1=rect.x1+1)
                rect_right = rect.copy(x2=rect.x2-1)
                result.extend(self._shrink_candidates(rect_left, depth=depth-1))
                result.extend(self._shrink_candidates(rect_right, depth=depth-1))

            if rect.height > 1:
                rect_top = rect.copy(y1=rect.y1+1)
                rect_bottom = rect.copy(y2=rect.y2-1)
                result.extend(self._shrink_candidates(rect_top, depth=depth-1))
                result.extend(self._shrink_candidates(rect_bottom, depth=depth-1))

        return result

    def _rects_reverse_projection(self, rects, img_shape, img_pad_shape, pad_top, pad_right, pad_bottom, pad_left):
        """Input images into the model are padded to make them squared. They
        are also resized to a smaller size. This function is supposed to
        remove both effects, i.e. to project the found bounding boxes from
        the padded and resized image to the unpadded und unresized (original)
        input image.
        """
        result = []
        for (rect, score) in rects:
            # project from resized padded (squared) image to unresized one
            rect_large = rect.on(img_pad_shape)
            # move rectangles to remove paddings
            rect_large_unpadded = rect_large.shift(top=-pad_top, left=-pad_left)
            # positions of corners are now correct, so switch underlying shape
            rect_large_unpadded = rect_large_unpadded.copy(shape=img_shape)
            result.append((rect_large_unpadded, score))
        return result

    def _rect_to_score(self, rect, heatmap):
        """Compute a score for a given rectangle (i.e. the confidence value).
        Currently this is done via an average of the corresponding activations
        in the heatmap.
        """
        subheatmap = rect.extract_from_image(heatmap)
        if subheatmap.ndim == 2 and subheatmap.shape[0] > 0 and subheatmap.shape[1] > 0:
            return np.average(subheatmap)
        else:
            print("[WARN] Broken heatmap extracted for rectangle", rect)
            return 0

    # The following stuff is some old code to make use of all generated
    # heatmaps. Didn't work well in tests.
    """
    def _merge_rects_to_bbs(self, hm_idx_to_rects, conf_threshold, img_shape):
        rects_full_size = self._make_rects_full_size(hm_idx_to_rects, self.img.shape)
        groups = self._group_rects(rects_full_size)

        final_bbs = []
        #for label, rects in cluster.iteritems():
        for label in groups:
            rects = groups[label]
            score_avg = sum([score for (rect, score) in rects]) / (1+9)
            if score_avg > conf_threshold:
                x1 = np.average([rect.x1 for (rect, score) in rects])
                x2 = np.average([rect.x2 for (rect, score) in rects])
                y1 = np.average([rect.y1 for (rect, score) in rects])
                y2 = np.average([rect.y2 for (rect, score) in rects])
                final_bbs.append((RectangleOnImage(x1=x1, y1=y1, x2=x2, y2=y2, shape=rect.shape), score_avg))

        return final_bbs

    def _make_rects_full_size(self, hm_idx_to_rects, img_orig_shape, keep_grouping=False):
        rects_full_size = []

        if keep_grouping:
            group = []
            rects_full_size.append(group)
        else:
            group = rects_full_size

        for (rect, score) in hm_idx_to_rects[0]:
            group.append((rect, score))

        nb_cells_y = 3
        nb_cells_x = 3
        grid_idx = 1
        for row_idx in range(nb_cells_y):
            for col_idx in range(nb_cells_x):
                if keep_grouping:
                    group = []
                    rects_full_size.append(group)
                else:
                    group = rects_full_size

                left = col_idx
                right = nb_cells_x - col_idx - 1
                above = row_idx
                below = nb_cells_y - row_idx - 1

                for (rect, score) in hm_idx_to_rects[grid_idx]:
                    x1 = rect.x1 - (left * rect.width)
                    x2 = rect.x2 + (right * rect.width)
                    y1 = rect.y1 - (above * rect.height)
                    y2 = rect.y2 + (below * rect.height)
                    rect_full_size = RectangleOnImage(x1=x1, x2=x2, y1=y1, y2=y2, shape=img_orig_shape)
                    group.append((rect_full_size, score))

                grid_idx += 1
        return rects_full_size

    def _group_rects(self, rects_full_size):
        if len(rects_full_size) == 0:
            return dict()
        elif len(rects_full_size) == 1:
            return dict([(0, [rects_full_size])])
        else:
            distances = np.zeros((len(rects_full_size), len(rects_full_size)), dtype=np.float32)
            for i in range(len(rects_full_size)):
                rect1 = rects_full_size[i][0]
                for j in range(i+1, len(rects_full_size)):
                    rect2 = rects_full_size[j][0]
                    sim = rect1.iou(rect2)
                    distances[i, j] = (1 - sim)
                    distances[j, i] = (1 - sim)

            clusterer = DBSCAN(metric="precomputed")
            labels = clusterer.fit_predict(distances)

            clusters = defaultdict(list)
            for label, (rect, score) in zip(labels, rects_full_size):
                clusters[label].append((rect, score))
            return clusters
    """

def debug_frame(frame_idx, img, model, conf_threshold, input_size=224):
    """Corresponding function to process_frame() that effectively does the same,
    but shows some debug information.

    Probably doesn't work currently as some functions were moved into a class.
    """
    img_orig_shape = img.shape
    img_pad, (pad_top, pad_right, pad_bottom, pad_left) = to_aspect_ratio_add(img, 1.0, return_paddings=True)
    img_rs = misc.imresize(img_pad, (input_size, input_size))
    inputs = (np.array([img_rs])/255.0).astype(np.float32).transpose(0, 3, 1, 2)
    inputs = torch.from_numpy(inputs)
    inputs = Variable(inputs)
    if GPU >= 0:
        inputs = inputs.cuda(GPU)
    outputs_pred = model(inputs)
    outputs_pred = outputs_pred.data.cpu().numpy()
    print("outputs_pred", np.min(outputs_pred), np.average(outputs_pred), np.max(outputs_pred))
    hm_idx_to_rects_pad = []
    for i in range(outputs_pred.shape[1]):
        hms = heatmap_to_rects(outputs_pred[0, i, ...], img_rs)
        hm_idx_to_rects_pad.append(hms)
    hm_idx_to_rects = []
    for i in range(outputs_pred.shape[1]):
        hms = heatmap_to_rects(outputs_pred[0, i, ...], img_rs)
        hms_rev = rects_reverse_projection(hms, img.shape, img_pad.shape, pad_top, pad_right, pad_bottom, pad_left)
        hm_idx_to_rects.append(hms_rev)

    rects_full_size_vis = make_rects_full_size(hm_idx_to_rects, img_orig_shape, keep_grouping=True)
    rects_full_size = make_rects_full_size(hm_idx_to_rects, img_orig_shape, keep_grouping=False)
    groups = group_rects(rects_full_size)

    final_bbs = []
    #for label, rects in cluster.iteritems():
    for label in groups:
        rects = groups[label]
        score_avg = sum([score for (rect, score) in rects]) / (1+9)
        if score_avg > conf_threshold:
            x1 = np.average([rect.x1 for (rect, score) in rects])
            x2 = np.average([rect.x2 for (rect, score) in rects])
            y1 = np.average([rect.y1 for (rect, score) in rects])
            y2 = np.average([rect.y2 for (rect, score) in rects])
            final_bbs.append((RectangleOnImage(x1=x1, y1=y1, x2=x2, y2=y2, shape=img_orig_shape), score_avg))

    img_rs_nopad = imresize_sidelen(img, 200, pick_func=max)
    rows = []

    # heatmaps
    row = [misc.imresize(img_rs, (img_rs_nopad.shape[0], img_rs_nopad.shape[1]))]
    for i in range(outputs_pred.shape[1]):
        hm = draw_heatmap(img_rs, outputs_pred[0, i])
        row.append(misc.imresize(hm, (img_rs_nopad.shape[0], img_rs_nopad.shape[1])))
    rows.append(np.hstack(row))

    # heatmaps => rects (padded image)
    #print("pad", pad_top, pad_right, pad_bottom, pad_left)
    #print("hm_idx_to_rects_pad", hm_idx_to_rects_pad)
    row = [misc.imresize(img_rs, (img_rs_nopad.shape[0], img_rs_nopad.shape[1]))]
    for rects in hm_idx_to_rects_pad:
        img_cp = np.copy(row[0])
        for (rect, score) in rects:
            img_cp = rect.draw_on_image(img_cp, color=[0, 255, 0])
        row.append(img_cp)
    rows.append(np.hstack(row))

    # heatmaps => rects (unpadded/original image)
    #print("hm_idx_to_rects", [[(r.on(img_rs_nopad), s) for (r, s) in rects] for rects in hm_idx_to_rects])
    row = [img_rs_nopad]
    for rects in hm_idx_to_rects:
        img_cp = np.copy(img_rs_nopad)
        for (rect, score) in rects:
            img_cp = rect.draw_on_image(img_cp, color=[0, 255, 0])
        row.append(img_cp)
    rows.append(np.hstack(row))

    # heatmaps => rects full size
    row = [img_rs_nopad]
    for rects in rects_full_size_vis:
        img_cp = np.copy(img_rs_nopad)
        for (rect, score) in rects:
            img_cp = rect.draw_on_image(img_cp, color=[0, 255, 0])
        row.append(img_cp)
    rows.append(np.hstack(row))

    # clustered rects
    img_cp = np.copy(img_rs_nopad)
    for label in groups:
        col = np.random.randint(0, 255, size=(3,))
        rects = groups[label]
        for (rect, score) in rects:
            img_cp = rect.draw_on_image(img_cp, color=col)
    row = np.hstack([img_rs_nopad, img_cp])
    diff = img_rs_nopad.shape[1] * (1+1+9) - row.shape[1]
    row = np.pad(row, ((0, 0), (0, diff), (0, 0)), mode="constant", constant_values=0)
    rows.append(row)

    # final rects
    img_cp = np.copy(img_rs_nopad)
    for (rect, score) in final_bbs:
        col = np.random.randint(0, 255, size=(3,))
        img_cp = rect.draw_on_image(img_cp, color=col)
    row = np.hstack([img_rs_nopad, img_cp])
    diff = img_rs_nopad.shape[1] * (1+1+9) - row.shape[1]
    row = np.pad(row, ((0, 0), (0, diff), (0, 0)), mode="constant", constant_values=0)
    rows.append(row)

    #print([r.shape for r in rows])
    misc.imshow(np.vstack(rows))

if __name__ == "__main__":
    main()
