# ckwg +28
# Copyright 2018-2019 by Kitware, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#  * Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
#  * Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
#  * Neither name of Kitware, Inc. nor the names of any contributors may be used
#    to endorse or promote products derived from this software without specific
#    prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS''
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from __future__ import print_function
from __future__ import division
from __future__ import absolute_import

import torch
import torch.nn as nn
import torch.utils.data as data

from torchvision import models, transforms, datasets

from PIL import Image as pilImage

from .parse_gpu_list import get_device
from kwiver.vital.types import BoundingBox

class AlexNetDataLoader( data.Dataset ):
    def __init__( self, bbox_list, transform, frame_img, in_size ):
        self._frame_img = pilImage.new( "RGB", frame_img.size )
        self._frame_img.paste( frame_img )
        self._transform = transform
        self._bbox_list = bbox_list
        self._in_size = in_size

    def __getitem__( self, index ):
        bb = self._bbox_list[index].bounding_box()

        # unwrap
        min_x = float( bb.min_x() )
        min_y = float( bb.min_y() )
        max_x = float( bb.max_x() )
        max_y = float( bb.max_y() )

        c_x = ( min_x + max_x ) / 2
        c_y = ( min_y + max_y ) / 2

        padding = 1.12 * max( max_x - min_x, max_y - min_y ) / 2

        # crop a square image
        im = self._frame_img.crop(
            (
                c_x - padding,
                c_y - padding,
                c_x + padding,
                c_y + padding
            )
        )

        im = im.resize((self._in_size, self._in_size), pilImage.BILINEAR)
        im.convert('RGB')

        if self._transform is not None:
            im = self._transform(im)

        return im

    def __len__(self):
        return self._bbox_list.size()


class AlexNetFeatureExtractor(object):
    """
    Obtain the appearance features from a trained pytorch alexnet
    model
    """

    def __init__(self, alexnet_model_path, img_size, batch_size, GPU_list=None):
        self._device, use_gpu_flag = get_device(GPU_list)

        # load the alexnet model. Maybe this shouldn't be hardcoded?
        self._alexnet_model = models.alexnet()
        print( alexnet_model_path )
        weights = torch.load( alexnet_model_path )

        self._alexnet_model.load_state_dict( weights )
        new_classifier = nn.Sequential(*list(self._alexnet_model.classifier.children())[:-1])
        self._alexnet_model.classifier = new_classifier

        self._alexnet_model.train( False ) # is this the same as eval() ?
        self._alexnet_model.to(self._device) # move the model to the GPU

        self._transform = transforms.Compose([
            transforms.Scale(img_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self._img_size = img_size
        self._b_size = batch_size

    def __call__(self, bbox_list, mot_flag):
        return self._obtain_feature(bbox_list, mot_flag)

    def _obtain_feature(self, bbox_list, mot_flag):
        kwargs = {'num_workers': 0, 'pin_memory': True}
        bbox_loader_class = AlexNetDataLoader(bbox_list, self._transform, self.frame, self._img_size)
        bbox_loader = torch.utils.data.DataLoader(bbox_loader_class, batch_size=self._b_size, shuffle=False, **kwargs)

        torch.set_grad_enabled(False)
        for idx, imgs in enumerate(bbox_loader):
            v_imgs = imgs.to(self._device)
            output = self._alexnet_model(v_imgs)

            if idx == 0:
                app_features = output.data
            else:
                app_features = torch.cat((app_features, output.data), dim=0)
        return app_features.cpu()
