# (╯‵□′)╯︵┻━┻

# 画面分割

import math
import cv2
import numpy as np
import numpy.linalg as npl

tileBackColor = (220, 223, 223)
tileBackColorInf = (215, 218, 218)
tileSup = (225, 228, 228)

horScanY = 640

# verbose print
verbose = False
def vprint(*o):
    if (verbose):
        print(*o)

def ptint(pt):
    return (int(pt[0]), int(pt[1]))

def v(*args):
    return np.array(args, dtype=float)

def vint(*args):
    return np.array(args, dtype=int)

def v32(*args):
    return np.array(args, dtype=np.float32)

# https://gist.github.com/meyerjo/dd3533edc97c81258898f60d8978eddc
def calcIou(boxA, boxB):
    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    # compute the area of intersection rectangle
    interArea = abs(max((xB - xA, 0)) * max((yB - yA), 0))
    if interArea == 0:
        return 0
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = abs((boxA[2] - boxA[0]) * (boxA[3] - boxA[1]))
    boxBArea = abs((boxB[2] - boxB[0]) * (boxB[3] - boxB[1]))

    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)

    # return the intersection over union value
    return iou

def colorSim(c1, c2, thr):
    return npl.norm(c1[:3] - c2[:3]) < thr

def colorLess(c1, c2):
    return (c1[0] < c2[0] and c1[1] < c2[1] and c1[2] < c2[2])

# --------------------------------手牌--------------------------------------

# 牌之间的缝
def isTileGap(img, x, u, d, tol=10):
    for y in range(u, d, 3):
        pixel = img[y, x]
        if (isTileBack(pixel, tol)):
            return False
    return True

# 牌底色
def isTileBack(c, tol=10):
    return colorSim(c, tileBackColor, 10) or colorLess(tileSup, c)

# 定位最左牌
def findLeftMargin(img):
    res = None
    for x in range(140, 180):
        pixel = img[horScanY, x]
        if (colorSim(pixel, tileBackColor, 10)):
            res = x
            break
    if (res == None):
        print('failed to find left margin')
        raise Exception('can not find left margin')
    return res

# 定位整张牌左右边界
def findTileHor(img, startX, u, d):
    l = None
    r = None
    # find left border
    hasBeenIn = False
    for x in range(startX + 30, startX - 10, -1):
        pixel = img[horScanY, x]
        if ((not isTileBack(pixel)) and isTileGap(img, x, u, d)):
            if (hasBeenIn):
                l = x + 1
                break
        else:
            hasBeenIn = True
    # find right border
    hasBeenIn = False
    if (l != None):
        for x in range(l + 55, l + 75, 1):
            pixel = img[horScanY, x]
            if ((not isTileBack(pixel)) and isTileGap(img, x, u, d)):
                if (hasBeenIn):
                    r = x
                    break
            else:
                hasBeenIn = True
    if (l == None or r == None):
        vprint('failed to find l or r (%s, %s)' % (l, r))
        return None
    return (l, r)

# 定位牌上下边界
def findVerticalBorder(img, x):
    u = None
    d = None
    for y in range(570, 640):
        pixel = img[y, x]
        if (isTileBack(pixel)):
            u = y
            break
    for y in range(719, 650, -1):
        pixel = img[y, x]
        if (isTileBack(pixel)):
            d = y
            break
    if (u == None or d == None):
        print('failed to find u or d (%s, %s)' % (u, d))
        raise Exception('can not find vertical border')
    return (u, d)

# 从全图提取自己的牌
def extractTilesImg(img):
    r = findLeftMargin(img)
    u, d = findVerticalBorder(img, r + 30)
    vprint('[extractTiles] left=%s u=%s d=%s' % (r, u, d))
    res = []
    while (True):
        lr = findTileHor(img, r, u, d)
        if (lr == None):
            vprint('[extractTiles] No more tiles')
            break

        l, r = lr
        u, d = findVerticalBorder(img, int((l + r) / 2))
        vprint('[extractTiles] new tile %s ~ %s  %s ~ %s' % (l, r, u, d))

        res.append((vint(l, u), vint(r, d), img[u:d, l:r]))
    if (len(res) == 0):
        raise Exception('no tile found!')
    return res

def wrapPers(m, xy):
    uv = np.matmul(m, [xy[0], xy[1], 1])
    uv /= uv[2]
    return uv[:2]

# --------------------------------牌河-------------------------------------

# 中间区域透视校正
def extractCenterRegeon(img):
    ulc = [349, 81]
    urc = [932, 81]
    blc = [237, 564]
    brc = [1043, 564]
    size = 800
    mat = cv2.getPerspectiveTransform(v32(ulc, urc, blc, brc),
        v32([0, 0], [size, 0], [0, size], [size, size]))
    res = cv2.warpPerspective(img, mat, (size, size))
    return res

# 底色分割
def tileBackRange(img):
    tol = 5
    #print(img.shape, (tileBackColor[0] - tol, tileBackColor[1] - tol, tileBackColor[2] - tol), (tileBackColor[0] + tol, tileBackColor[1] + tol, tileBackColor[2] + tol))
    seg = cv2.inRange(img,
        (tileBackColor[0] - tol, tileBackColor[1] - tol, tileBackColor[2] - tol),
        (tileBackColor[0] + tol, tileBackColor[1] + tol, tileBackColor[2] + tol))
    seg2 = cv2.inRange(img, tileSup, (255, 255, 255))
    res = cv2.bitwise_or(seg, seg2)
    #res = cv2.Laplacian(res, cv2.CV_8U, ksize=3)
    #res = cv2.GaussianBlur(res, (5, 5), 3)
    return res

## 双向不等价的角度比较
#def angleSim(a1, a2, thr = 10 / 180 * np.pi):
#    return abs((a1 - a2 + np.pi) % (2 * np.pi) - np.pi) < thr
#
## 双向等价的角度比较
#def angleSim2(a1, a2, thr = 10 / 180.0 * np.pi):
#    return abs((a1 - a2 + np.pi/2) % np.pi - np.pi/2) < thr
#
## check vertical edge at x from ys to ye
#def checkVerticalEdge(img_mag, img_dir, x, ys, ye):
#    for y in range(ys, ye+1, 2):
#        px_mag = img_mag[y, x]
#        px_dir = img_dir[y, x]
#        if (not angleSim2(px_dir, 0)):
#            return False
#    return True
#
## check if edge exists
#def checkEdge(img_mag, img_dir, pstart, pend):
#    pstart = np.array(pstart, dtype=float)
#    pend = np.array(pend, dtype=float)
#    edge_dir = (pend - pstart)
#    edge_len = npl.norm(edge_dir)
#    edge_dir /= edge_len
#    edge_normdir = math.atan2(-edge_dir[0], edge_dir[1])
#    for i in range(0, math.ceil(edge_len), 2):
#        x, y = ptint(pstart + i * edge_dir)
#        px_dir = img_dir[y, x]
#        px_mag = img_mag[y, x]
#        if ((angleSim2(px_dir, edge_normdir)) or (px_mag < 50)):
#            pass # safe
#        else:
#            return False
#    return True

# get match box with best iou   return box,iou
def getBestMatchBox(box, boxlist):
    bestBox = None
    bestIou = 0
    for box1 in boxlist:
        iou = calcIou([box[0][0], box[0][1], box[1][0], box[1][1]],
            [box1[0][0], box1[0][1], box1[1][0], box1[1][1]])
        if (iou > bestIou):
            bestIou = iou
            bestBox = box1
    return bestBox, bestIou # could be None

# check shape
def isInGoodShape(box, p_width, p_height):
    p1, p2 = box
    w = p2[0] - p1[0]
    h = p2[1] - p1[1]
    area = w * h
    p_area = p_width * p_height
    if (max(p_area/area, area/p_area) > 1.1):
        return False
    if (max(w/p_width, p_width/w) > 1.05):
        return False
    if (max(h/p_height, p_height/h) > 1.05):
        return False
    return True

# 一家的打出的牌
def extractPartTilesImg(img):
    img_seg = tileBackRange(img)
    edged = cv2.Canny(img_seg, 30, 200)
    contours, hierarchy = cv2.findContours(edged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    vprint('[extractPartTilesImg] found %s contours' % len(contours))
    contours_outer = []
    for i in range(len(contours)):
        if (hierarchy[0][i][3] == -1):
            contours_outer.append(contours[i])

    # 粗筛
    #canvas = img.copy()
    rect_candidates = []
    for contour in contours_outer:
        x,y,w,h = cv2.boundingRect(contour)
        if (w * h >= 600 and max(w/h, h/w) < 4):
            rect_candidates.append(([x, y], [x+w, y+h]))
            #cv2.rectangle(canvas, [x, y], [x+w, y+h], (0, 255, 0), 1)
        else:
            #cv2.rectangle(canvas, [x, y], [x+w, y+h], (0, 255, 255), 1)
            pass
    vprint('[extractPartTilesImg] %s candidates' % len(rect_candidates))

    # prior params
    p_width = 43 # 牌宽
    p_height = 56 # 牌高
    p_winterval = 46 # 竖直摆放的水平间隔
    p_wwinterval = 61 # 水平摆放的水平间隔
    p_hinterval = 62 # 竖直摆放的竖直间隔
    p_horsink = 6 # 水平摆放的高度差
    p_start = [10, 10]
    p_maxX = 260 # maximum line length
    p_maxLine = 3 # maximum line num
    # match window
    cur = p_start
    curLine = 0
    rects = []
    while (True):
        estBox= [cur, [cur[0]+p_width, cur[1]+p_height]]
        estBoxHor = [[cur[0], cur[1]+p_horsink], [cur[0]+p_height, cur[1]+p_width]]
        vprint('[extractPartTilesImg] check est box %s %s' % (estBox, estBoxHor))
        bestBox, iou = getBestMatchBox(estBox, rect_candidates)
        bestBoxHor, iouHor = getBestMatchBox(estBoxHor, rect_candidates)
        vprint('[extractPartTilesImg] matching box %s(%s) %s(%s)' % (bestBox, iou, bestBoxHor, iouHor))
        if (bestBox == None):
            if (cur[0] > p_maxX):
                # try next line
                vprint('[extractPartTilesImg] line break')
                curLine += 1
                if (curLine >= p_maxLine):
                    #oh no
                    break
                cur = [p_start[0], p_start[1] + p_hinterval * curLine]
            else:
                # maybe try next
                vprint('[extractPartTilesImg] no luck, skipping tile')
                cur = [cur[0] + p_winterval, cur[1]]
        else:
            # something is here
            if (iou > (iouHor-0.1)): # more likely to be vertical
                #it might be vertical
                vprint('[extractPartTilesImg] found vertical tile')
                if (isInGoodShape(bestBox, p_width, p_height)):
                    rects.append((bestBox, iou))
                    cur = [bestBox[0][0] + p_winterval, bestBox[0][1]]
                else:
                    rects.append((estBox, iou))
                    cur = [estBox[0][0] + p_winterval, estBox[0][1]]
            else:
                #it might be horizontal
                vprint('[extractPartTilesImg] found horizontal tile')
                if (isInGoodShape(bestBoxHor, p_height, p_width)):
                    rects.append((bestBoxHor, iouHor))
                    cur = [bestBox[0][0] + p_wwinterval, bestBox[0][1] - p_horsink]
                else:
                    rects.append((estBoxHor, iouHor))
                    cur = [estBoxHor[0][0] + p_wwinterval, estBoxHor[0][1] - p_horsink]

#    for p1, p2 in rects:
#        cv2.rectangle(canvas, p1, p2, (0, 0, 255), 1)
#    cv2.imshow('canvas', canvas)
#    cv2.waitKey(0)

    res = []
    for (p1, p2), iou in rects:
        w = p2[0] - p1[0]
        h = p2[1] - p1[1]
        if (w > h):
            res.append((vint(*p1), vint(*p2), cv2.rotate(img[p1[1]:p2[1], p1[0]:p2[0]], cv2.ROTATE_90_COUNTERCLOCKWISE)))
        else:
            res.append((vint(*p1), vint(*p2), img[p1[1]:p2[1], p1[0]:p2[0]]))
#        cv2.imshow('prev', res[-1])
#        cv2.waitKey()

    return res

# 从全图提取
def extractCenterTilesImg(img):
    # extract center region
    center = extractCenterRegeon(img)
    blcorner = (252, 523)
    brcorner = (542, 525)
    urcorner = (545, 241)
    ulcorner = (256, 231)
    secH = 210
    secH = [210, 210, 210, 210]
    secW = [290, 290, 290, 290]
    parts = [None]*4
    parts[0] = center[blcorner[1]:blcorner[1]+secH[0], blcorner[0]:blcorner[0]+secW[0]].copy()
    parts[1] = cv2.rotate(center[brcorner[1]-secW[1]:brcorner[1], brcorner[0]:brcorner[0]+secH[1]], cv2.ROTATE_90_CLOCKWISE)
    parts[2] = cv2.rotate(center[urcorner[1]-secH[2]:urcorner[1], urcorner[0]-secW[2]:urcorner[0]], cv2.ROTATE_180)
    parts[3] = cv2.rotate(center[ulcorner[1]:ulcorner[1]+secW[3], ulcorner[0]-secH[3]:ulcorner[0]], cv2.ROTATE_90_COUNTERCLOCKWISE)

    tiles = []
    for i, part in enumerate(parts):
        vprint('[extractCenterTilesImg] part %s' % i)
        tiles.append(extractPartTilesImg(part))
    return tiles

# --------------------------------其他杂项--------------------------------------

def testRange(img):
    img = img[353:410, 434:567]

    num_pixel = img.shape[0] * img.shape[1]
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    kernel = np.ones((5,5),np.uint8)
    chii_range = cv2.morphologyEx(cv2.inRange(hsv, (75, 50, 50), (83, 255, 255)), cv2.MORPH_OPEN, kernel)
    pon_range = cv2.morphologyEx(cv2.inRange(hsv, (97, 50, 50), (100, 255, 255)), cv2.MORPH_OPEN, kernel)
    kan_range = cv2.morphologyEx(cv2.inRange(hsv, (150, 50, 50), (155, 255, 255)), cv2.MORPH_OPEN, kernel)
    ron_range = cv2.morphologyEx(cv2.bitwise_or(
        cv2.inRange(hsv, (0, 50, 50), (2, 255, 255)),
        cv2.inRange(hsv, (178, 50, 50), (180, 255, 255))  ), cv2.MORPH_OPEN, kernel)
    riichi_range = cv2.morphologyEx(cv2.inRange(hsv, (10, 50, 50), (14, 255, 255)), cv2.MORPH_OPEN, kernel)

    chii_ratio = np.sum(chii_range) / 255.0 / num_pixel
    pon_ratio = np.sum(pon_range) / 255.0 / num_pixel
    kan_ratio = np.sum(kan_range) / 255.0 / num_pixel
    ron_ratio = np.sum(ron_range) / 255.0 / num_pixel
    riichi_ratio = np.sum(riichi_range) / 255.0 / num_pixel

    cv2.imshow('chii', chii_range)
    cv2.imshow('pon', pon_range)
    cv2.imshow('kan', kan_range)
    cv2.imshow('ron', ron_range)
    cv2.imshow('riichi', riichi_range)

    print(chii_ratio, pon_ratio, kan_ratio, ron_ratio, riichi_ratio)

    cv2.waitKey()

# 碰吃杠立直和
def extractActions(img):
    actionROI = img[527:527+58, 614:614+164]

    num_pixel = actionROI.shape[0] * actionROI.shape[1]
    hsv = cv2.cvtColor(actionROI, cv2.COLOR_BGR2HSV)

    kernel = np.ones((5,5),np.uint8)
    chii_range = cv2.morphologyEx(cv2.inRange(hsv, (75, 50, 50), (83, 255, 255)), cv2.MORPH_OPEN, kernel)
    pon_range = cv2.morphologyEx(cv2.inRange(hsv, (97, 50, 50), (100, 255, 255)), cv2.MORPH_OPEN, kernel)
    kan_range = cv2.morphologyEx(cv2.inRange(hsv, (150, 50, 50), (155, 255, 255)), cv2.MORPH_OPEN, kernel)
    ron_range = cv2.morphologyEx(cv2.bitwise_or(
        cv2.inRange(hsv, (0, 50, 50), (2, 255, 255)),
        cv2.inRange(hsv, (178, 50, 50), (180, 255, 255))  ), cv2.MORPH_OPEN, kernel)
    riichi_range = cv2.morphologyEx(cv2.inRange(hsv, (10, 50, 50), (14, 255, 255)), cv2.MORPH_OPEN, kernel)
    tsumo_range = cv2.morphologyEx(cv2.inRange(hsv, (124, 50, 50), (127, 255, 255)), cv2.MORPH_OPEN, kernel)

    action_names = ['chii', 'pon', 'kan', 'ron', 'riichi', 'tsumo']
    action_ratio = [
        np.sum(chii_range) / 255.0 / num_pixel,
        np.sum(pon_range) / 255.0 / num_pixel,
        np.sum(kan_range) / 255.0 / num_pixel,
        np.sum(ron_range) / 255.0 / num_pixel,
        np.sum(riichi_range) / 255.0 / num_pixel,
        np.sum(tsumo_range) / 255.0 / num_pixel
    ]
    print(action_ratio)
    action_idx = np.argmax(action_ratio)
    if (action_ratio[action_idx] < 0.06):
        return None
    return action_names[action_idx]


# -*- coding: utf-8 -*-
'''VGG16-places365 model for Keras

# Reference:
- [Places: A 10 million Image Database for Scene Recognition](http://places2.csail.mit.edu/PAMI_places.pdf)
'''

from __future__ import division, print_function
import os

import warnings
import numpy as np

from keras import backend as K
from keras.layers import Input
from keras.layers.core import Activation, Dense, Flatten
from keras.layers.pooling import MaxPooling2D
from keras.models import Model
from keras.layers import Conv2D
from keras.regularizers import l2
from keras.layers.core import Dropout
from keras.layers import GlobalAveragePooling2D
from keras.layers import GlobalMaxPooling2D
from keras_applications.imagenet_utils import _obtain_input_shape
from keras.engine.topology import get_source_inputs
from keras.utils.data_utils import get_file
from keras.utils import layer_utils
from keras.preprocessing import image
from keras_applications.imagenet_utils import preprocess_input

WEIGHTS_PATH = 'https://github.com/GKalliatakis/Keras-VGG16-places365/releases/download/v1.0/vgg16-places365_weights_tf_dim_ordering_tf_kernels.h5'
WEIGHTS_PATH_NO_TOP = 'https://github.com/GKalliatakis/Keras-VGG16-places365/releases/download/v1.0/vgg16-places365_weights_tf_dim_ordering_tf_kernels_notop.h5'

def VGG16_Places365(include_top=True, weights='places',
                    input_tensor=None, input_shape=None,
                    pooling=None,
                    classes=365):
    """Instantiates the VGG16-places365 architecture.

    Optionally loads weights pre-trained
    on Places. Note that when using TensorFlow,
    for best performance you should set
    `image_data_format="channels_last"` in your Keras config
    at ~/.keras/keras.json.

    The model and the weights are compatible with both
    TensorFlow and Theano. The data format
    convention used by the model is the one
    specified in your Keras config file.

    # Arguments
        include_top: whether to include the 3 fully-connected
            layers at the top of the network.
        weights: one of `None` (random initialization),
                 'places' (pre-training on Places),
                 or the path to the weights file to be loaded.
        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
            to use as image input for the model.
        input_shape: optional shape tuple, only to be specified
            if `include_top` is False (otherwise the input shape
            has to be `(224, 224, 3)` (with `channels_last` data format)
            or `(3, 224, 244)` (with `channels_first` data format).
            It should have exactly 3 inputs channels,
            and width and height should be no smaller than 48.
            E.g. `(200, 200, 3)` would be one valid value.
        pooling: Optional pooling mode for feature extraction
            when `include_top` is `False`.
            - `None` means that the output of the model will be
                the 4D tensor output of the
                last convolutional layer.
            - `avg` means that global average pooling
                will be applied to the output of the
                last convolutional layer, and thus
                the output of the model will be a 2D tensor.
            - `max` means that global max pooling will
                be applied.
        classes: optional number of classes to classify images
            into, only to be specified if `include_top` is True, and
            if no `weights` argument is specified.
    # Returns
        A Keras model instance.
    # Raises
        ValueError: in case of invalid argument for `weights`, or invalid input shape
        """
    if not (weights in {'places', None} or os.path.exists(weights)):
        raise ValueError('The `weights` argument should be either '
                         '`None` (random initialization), `places` '
                         '(pre-training on Places), '
                         'or the path to the weights file to be loaded.')

    if weights == 'places' and include_top and classes != 365:
        raise ValueError('If using `weights` as places with `include_top`'
                         ' as true, `classes` should be 365')

    # Determine proper input shape
    input_shape = _obtain_input_shape(input_shape,
                                      default_size=224,
                                      min_size=48,
                                      data_format=K.image_data_format(),
                                      require_flatten =include_top)

    if input_tensor is None:
        img_input = Input(shape=input_shape)
    else:
        if not K.is_keras_tensor(input_tensor):
            img_input = Input(tensor=input_tensor, shape=input_shape)
        else:
            img_input = input_tensor

    # Block 1
    x = Conv2D(filters=64, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block1_conv1')(img_input)

    x = Conv2D(filters=64, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block1_conv2')(x)

    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), name="block1_pool", padding='valid')(x)

    # Block 2
    x = Conv2D(filters=128, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block2_conv1')(x)

    x = Conv2D(filters=128, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block2_conv2')(x)

    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), name="block2_pool", padding='valid')(x)

    # Block 3
    x = Conv2D(filters=256, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block3_conv1')(x)

    x = Conv2D(filters=256, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block3_conv2')(x)

    x = Conv2D(filters=256, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block3_conv3')(x)

    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), name="block3_pool", padding='valid')(x)

    # Block 4
    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block4_conv1')(x)

    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block4_conv2')(x)

    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block4_conv3')(x)

    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), name="block4_pool", padding='valid')(x)

    # Block 5
    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block5_conv1')(x)

    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block5_conv2')(x)

    x = Conv2D(filters=512, kernel_size=3, strides=(1, 1), padding='same',
               kernel_regularizer=l2(0.0002),
               activation='relu', name='block5_conv3')(x)

    x = MaxPooling2D(pool_size=(2, 2), strides=(2, 2), name="block5_pool", padding='valid')(x)

    if include_top:
        # Classification block
        x = Flatten(name='flatten')(x)
        x = Dense(4096, activation='relu', name='fc1')(x)
        x = Dropout(0.5, name='drop_fc1')(x)

        x = Dense(4096, activation='relu', name='fc2')(x)
        x = Dropout(0.5, name='drop_fc2')(x)

        x = Dense(365, activation='softmax', name="predictions")(x)

    else:
        if pooling == 'avg':
            x = GlobalAveragePooling2D()(x)
        elif pooling == 'max':
            x = GlobalMaxPooling2D()(x)

    # Ensure that the model takes into account
    # any potential predecessors of `input_tensor`.
    if input_tensor is not None:
        inputs = get_source_inputs(input_tensor)
    else:
        inputs = img_input

    # Create model.
    model = Model(inputs, x, name='vgg16-places365')

    # load weights
    if weights == 'places':
        if include_top:
            weights_path = get_file('vgg16-places365_weights_tf_dim_ordering_tf_kernels.h5',
                                    WEIGHTS_PATH,
                                    cache_subdir='models')
        else:
            weights_path = get_file('vgg16-places365_weights_tf_dim_ordering_tf_kernels_notop.h5',
                                    WEIGHTS_PATH_NO_TOP,
                                    cache_subdir='models')

        model.load_weights(weights_path)

        if K.backend() == 'theano':
            layer_utils.convert_all_kernels_in_model(model)

        if K.image_data_format() == 'channels_first':
            if include_top:
                maxpool = model.get_layer(name='block5_pool')
                shape = maxpool.output_shape[1:]
                dense = model.get_layer(name='fc1')
                layer_utils.convert_dense_weights_data_format(dense, shape, 'channels_first')

            if K.backend() == 'tensorflow':
                warnings.warn('You are using the TensorFlow backend, yet you '
                              'are using the Theano '
                              'image data format convention '
                              '(`image_data_format="channels_first"`). '
                              'For best performance, set '
                              '`image_data_format="channels_last"` in '
                              'your Keras config '
                              'at ~/.keras/keras.json.')

    elif weights is not None:
        model.load_weights(weights)

    return model

if __name__ == '__main__':
    import urllib2
    import numpy as np
    from PIL import Image
    from cv2 import resize

    TEST_IMAGE_URL = 'http://places2.csail.mit.edu/imgs/demo/6.jpg'

    image = Image.open(urllib2.urlopen(TEST_IMAGE_URL))
    image = np.array(image, dtype=np.uint8)
    image = resize(image, (224, 224))
    image = np.expand_dims(image, 0)

    model = VGG16_Places365(weights='places')
    predictions_to_return = 5
    preds = model.predict(image)[0]
    top_preds = np.argsort(preds)[::-1][0:predictions_to_return]

    # load the class label
    file_name = 'categories_places365.txt'
    if not os.access(file_name, os.W_OK):
        synset_url = 'https://raw.githubusercontent.com/csailvision/places365/master/categories_places365.txt'
        os.system('wget ' + synset_url)
    classes = list()
    with open(file_name) as class_file:
        for line in class_file:
            classes.append(line.strip().split(' ')[0][3:])
    classes = tuple(classes)

    print('--PREDICTED SCENE CATEGORIES:')
    # output the prediction
    for i in range(0, 5):
        print(classes[top_preds[i]])

    # --PREDICTED SCENE CATEGORIES:
    # cafeteria
    # food_court
    # restaurant_patio
    # banquet_hall
    # restaurant


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: deep_ml_curriculum
#     language: python
#     name: deep_ml_curriculum
# ---

# # Time Series Anaysis
# Time Series are a sequence of data points ordered in time. Each data point contains one or more values and also has a time stamp. For instance, the average daily temperature is a time series. Time Series Analysis (TSA) is studying the behaviour of time series.
#
# An important part of time series analysis (TSA) is time series forcasting (TSF) which is using a sequence of data to predict future behaviour. In this session we will discuss TSA, and tool and technques in TSA. The next session will be focused on forcasting.
#
#
# TSA has its own unique vocabulary and technical terms. Throughout this notebook we will explain some of these terms and what they mean in practice.
#
#
# ## Tools
#
# One of the main tools for TSA is __Pandas__. In previous notebooks we learned about Pandas and it's capabilities in reading, writing, and manipulating tabular data. Pandas also has a wide range of functionalities for TSA.
#
# Another common tool for TSA in python is __Statsmodels__. Statsmodels is library with wide range of statistical analysis tools and a sub-module dedicated to TSA.

# +
# Our main packages for TSA
import pandas as pd

# For math and array operations
import numpy as np

# For plotting
import matplotlib.pyplot as plt
# -

# Later, we will import statsmodels sub-modules as we need them.

# ## Pandas for time series
#
# Pandas has functionalities for reading time and date.

data_normal = pd.read_csv("../../data/processed/Generated/SampleSales.csv")
data_normal["Date"]

# When reading data that contains date/time we can ask pandas to parse it as datetimes with `parse_dates = ['Column_name']`.

data_date_parsed = pd.read_csv(
    "../../data/processed/Generated/SampleSales.csv", parse_dates=["Date"]
)
data_date_parsed["Date"]

# In the first scenario when we print 'Date' column at the bottom it prints `dtype = object` which in this case it means pandas understands the dates as a string (or text). It has no understanding of what the dates actually mean. But in the second scenario, the type is `datetime64[ns]` which is a date/time object with accuracy of nanoseconds. Now pandas knows these are dates and it knows the year, month, day, etc. Also, despite showing only year, month, and day it is saving it to nanoseconds accuracy, which means we could also have time of the day included as well.
#

# Also, by passing `index_col = 'Date'` we can let pandas know that in the index column, instead of using numbers (0, 1, 2, ...) use the _'Date'_ column. This will make it easier to access data based on date.

data_index_parsed = pd.read_csv(
    "../../data/processed/Generated/SampleSales.csv",
    parse_dates=["Date"],
    index_col=["Date"],
)

# Notice freq is none, that makes sense since we don't have values for every day
data_index_parsed.index

data_index_parsed['Total'].plot()

data_index_parsed

# Now we can get the data for a certain date or time period. Let stick to the ISO format which is year-month-day, this sidesteps the confusion over US vs UK date formats.

data_index_parsed.loc["2014-12-19"]

data_index_parsed.loc["2014-11-01":"2014-12-01"]

# Also, if you have an array of dates in string format you can use `.to_datetime()` to convert it to datetime object. The date string can be in many formats and pandas can usually figure it out.

pd.to_datetime(["1/3/2019", "July 23rd, 1964", "2006-01-09"])

# Or you can specify which format pandas should use.

pd.to_datetime(["1/3/2019"], format="%d/%m/%Y")

pd.to_datetime(["1/3/2019"], format="%m/%d/%Y")

# ## Resampling
# There are times you might need weekly data but you have it on daily basis. You can use resampling to transform the data into the new time intervals. This is also useful when you are dealing with datapoints which are not equally spaced. Let's have another look at the data we loaded earlier.

data_index_parsed.head(10)

# This data shows the time and amount of purchases from a shop. The purchases are not taking place on a fixed interval. They can happen anytime. We can resample this data into weekly sales which makes it easier to analyse. To do this, we need to use the `.resample()` method and pass in the new interval. After that we need to specify how we want the data in each interval to be aggregated. In this case, we choose the sum function because we want the total sales in each week.

df = data_index_parsed.resample("W").sum()
df.head(10)

# Now we have weekly sales. Note that we passed in "W" to convert the data into weekly sales. There are a range of options available in pandas for various intervals. The table below shows a range of options for resampling.

# Note that not all the options on the table below are applicable for resampling. The options for resampling are "M", "A", "Q", "BM", "BA", "BQ", and "W". The other options in the table have other application which we will discuss later.

# DateOffsets and their frequency strings. ([Source](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#dateoffset-objects))

#
#

#

#

#       Date Offset
#       Frequency String
#       Description
#

#

#

#

#       DateOffset
#       None
#       Generic offset class, defaults to 1 calendar day
#

#

#       BDay or BusinessDay
#       'B'
#       business day (weekday)
#

#

#       CDay or CustomBusinessDay
#       'C'
#       custom business day
#

#

#       Week
#       'W'
#       one week, optionally anchored on a day of the week
#

#

#       WeekOfMonth
#       'WOM'
#       the x-th day of the y-th week of each month
#

#

#       LastWeekOfMonth
#       'LWOM'
#       the x-th day of the last week of each month
#

#

#       MonthEnd
#       'M'
#       calendar month end
#

#

#       MonthBegin
#       'MS'
#       calendar month begin
#

#

#       BMonthEnd or BusinessMonthEnd
#       'BM'
#       business month end
#

#

#       BMonthBegin or BusinessMonthBegin
#       'BMS'
#       business month begin
#

#

#       CBMonthEnd or CustomBusinessMonthEnd
#       'CBM'
#       custom business month end
#

#

#       CBMonthBegin or CustomBusinessMonthBegin
#       'CBMS'
#       custom business month begin
#

#

#       SemiMonthEnd
#       'SM'
#       15th (or other day_of_month) and calendar month end
#

#

#       SemiMonthBegin
#       'SMS'
#       15th (or other day_of_month) and calendar month begin
#

#

#       QuarterEnd
#       'Q'
#       calendar quarter end
#

#

#       QuarterBegin
#       'QS'
#       calendar quarter begin
#

#

#       BQuarterEnd
#       'BQ
#       business quarter end
#

#

#       BQuarterBegin
#       'BQS'
#       business quarter begin
#

#

#       FY5253Quarter
#       'REQ'
#       retail (aka 52-53 week) quarter
#

#

#       YearEnd
#       'A'
#       calendar year end
#

#

#       YearBegin
#       'AS' or 'BYS'
#       calendar year begin
#

#

#       BYearEnd
#       'BA'
#       business year end
#

#

#       BYearBegin
#       'BAS'
#       business year begin
#

#

#       FY5253
#       'RE'
#       retail (aka 52-53 week) year
#

#

#       Easter
#       None
#       Easter holiday
#

#

#       BusinessHour
#       'BH'
#       business hour
#

#

#       CustomBusinessHour
#       'CBH'
#       custom business hour
#

#

#       Day
#       'D'
#       one absolute day
#

#

#       Hour
#       'H'
#       one hour
#

#

#       Minute
#       'T' or 'min'
#       one minute
#

#

#       Second
#       'S'
#       one second
#

#

#       Milli
#       'L' or 'ms'
#       one millisecond
#

#

#       Micro
#       'U' or 'us'
#       one microsecond
#

#

#       Nano
#       'N'
#       one nanosecond
#

#

#

#

# Lets compare differen't resamplings
ax=plt.gca()
data_index_parsed['Total'].resample("D").sum().plot(label="Day", ax=ax)
data_index_parsed['Total'].resample("W").sum().plot(label="Week", ax=ax)
data_index_parsed['Total'].resample("M").sum().plot(label="Month", ax=ax)
plt.ylabel('Total')
plt.legend()

# ## Shifting
# Shifiting is moving the data forward or backward in time. For instance, the code below shifts the data one week to the future.

# Pandas knows the interval is a week
df.info()

df.head()

df.shift(periods=1)

# Notice the value that belonged to "2014-11-09" now belongs to "2014-11-09". Also, you can see in the first row since there was no data point to be moved there, its new value is `NaN`.

# Now the question is how does pandas know how much it should shift the data? In other words, why did it decide to move the data one week and not one day? The answer is because pandas knows our data is weekly, so when we asked it to shift the data by one period it assumed one week. This is called the frequency of the data and we can check it using `.info()` as we did above.

# The frequency is "W-SUN" which means weekly data and each week starts at sunday.

# __What if we want Monday to be start of the week?__

# When we were resampling the data we asked pandas to create weekly data. We can specify which day should start of the week.

data_index_parsed.resample("W-MON").sum()

# __What if we only wanted to shift the data one day and not one week?__

# We can specify a frequency when shifting.

df.shift(periods=1, freq="D")

# Note in the above example only the dates changed and `NaN` didn't appear in the first row.

# We can also shift backward. If we do this, `NaN` will appear at the end of the data.

df.shift(-2)

# ## Rolling
# Rolling is applying a function to a number of consecutive data points. For instance, at each point in time we take the average of the past 4 weeks. Taking average of past few periods is also know as __Moving Average Smoothing__.

# We first need to select the window (how many values we should consider), and then apply an aggregation function. (e.g. mean, sum, etc.)
#
#
#
#

#

df.rolling(window=4).mean()

# __Note:__ Since the first three points didn't have three instances before them, the moving average was not calculated for them.

# Let's save the moving average in the same data frame and plot the result to better understand why we would do this.

df["MA-4"] = df.rolling(window=4).mean()  # MA-4 for moving average of size 4 window
df.plot(y=["Total", "MA-4"], figsize=(16, 8))

# Let's try a larger window size. Because now we have two columns in the data, we need to specify which column we are taking a moving average of.

df["MA-13"] = (
    df["Total"].rolling(window=13).mean()
)  # MA-4 for moving average of size 4 window
df.plot(y=["Total", "MA-4", "MA-13"], figsize=(16, 8))

# As the name suggests __Moving Average Smoothing__ creates a smoothed out version of the data. The larger the window size, smoother the data will be. This is helpful when we want to get rid of noise in the data so we can have a better look at the trends. However, noise reduction comes at a cost. The downside of using moving average is the it creates lag in the data. For instance if there is a sudden increase in the data moving average responds to this change slower. It also doesn't reach the peaks of the original data. In the plot above you can see that "MA-13" does not show the high and low values of the data but only follows th main trend.

# Another useful function pandas provide is **expanding**. While at each point rolling applied a function over a fixed window, expanding applies a function over all the points that came before that specific point. For instance, if we apply mean function at 10th point we will have the average of first 10 points and at 20th point we will have the average of first 20 points.
#
# exanding().sum() is the same as .cumsum()
#

df["Total"].expanding().sum().plot()

# exanding().sum() is the same as .cumsum(), but you can use operations other than sum
df["Total"].cumsum().plot()

#

#
Exercise

#
# Here's an exercise to get you familiar with the Pandas docs on timeseries.
#
# You need to change the timezone on a dataframe but your not sure how.
#
# - Look through [the dataframe docs](https://pandas.pydata.org/docs/reference/frame.html) and find how to set a timezone on a series that doesn't have one (but you know it's 'Australia/Perth')
# - Set it to 'Australia/Perth'
# - plot it
#
#
#
#   → Hints
#
#   * Start in the Time-Series related section
#   * Scan it, then Ctrl-F for `timezone`
#   * Jargon is `localize` a `tz naive` series
#
#
#
#

#

#
#
#     → Solution
#
#
#   ```python
#   # DataFrame.tz_localize
#   df.tz_localize('Australia/Perth').plot(
#   ```
#
#
#
#

# Other operations:
#
# - https://pandas.pydata.org/pandas-docs/stable/reference/general_functions.html#top-level-dealing-with-datetimelike
# - https://pandas.pydata.org/pandas-docs/stable/reference/frame.html#time-series-related
# - https://pandas.pydata.org/pandas-docs/stable/reference/series.html#time-series-related

# ## Lowess Smoothing
# Now that we learned about moving average smoothing, let's discuss another technique. Locally weighted scatterplot smoothing (LOWESS) is another technique for smoothing.
#
# Lowess creates local models for small subsets of data and uses a weighted average of them for smoothing.
#

# The implementation of Lowess in statsmodels can be found in nonparametric sub-module.
from statsmodels.nonparametric.smoothers_lowess import lowess

# To create a smoothed output it requires x and y values. For time series x would be the time dimension and y would be the values.

x = df.index
y = df["Total"]
result = lowess(y, x)
df["Lowess"] = result[
    :, 1
]  # Lowess returns x and y. The second column (column index 1) are the values we need.
df.plot(y=["Total", "Lowess"])

# ## Time Series Decomposition

# The goal of decomposition is to break down time series into different elements each describing a specific behaviour of data.

# ### STL Decomposition

# Seasonal-Trend decomposition with LOESS (STL) is another method for decomposition of time series. This method breaks down the data into three elements:
# 1. Trend
# 2. Seasonality
# 3. Residual (Error or Noise)
#
# We learned about LOWESS for smoothing (LOESS is a similar method with only minor differences). LOESS showed the capability to extract the trend of time series. By removing the trend from the data STL assumes the remainder is sum of a seasonal component and some residual. Seasonal component is the repeating pattern of data. The sort of behaviour that repeats weekly, monthly, etc.
STL can decompose the time serie in two ways. First is additive method, where the original data is sum of its components:
# $$y_t = \tau_t + s_t + r_t $$
# Second, the multiplicative approach where the data is product of its components:
# $$y_t = \tau_t \times s_t \times r_t $$
# Statsmodel has a function for STL called `seasonal_decompose`.

# This data shows the number of customers a retail shop had per months
cust = pd.read_csv(
    "../../data/processed/Generated/Customers.csv", parse_dates=True, index_col="Month"
)
cust.head()

# Let's have quick look at the data.

cust.plot(y="Customers")

# We can clearly see a repeating pattern. That's called seasonality and we expect STL to be able to capture this behaviour.

from statsmodels.tsa.seasonal import seasonal_decompose
result = seasonal_decompose(cust["Customers"], model='additive')

# Now the time serie is decomposed and we can plot the result.

ax = result.plot()
ax.set_figwidth(16)
ax.set_figheight(8)

# By default the model used additive approach for the data. If we want to use multiplicative approach we need to specify it in the function. But what is the difference between additive and multiplicative? Additive approach is used when the height of the repeating patterns are constant. This means the amplitude of seasonal component is not changing. But in this data we can clearly see the seasonal component is getting larger and larger. This is also visible in the residual component. The residual component is meant to look like random noise but we can clearly see some regular oscillation in there. So let's try multiplicative approach.

result = seasonal_decompose(cust["Customers"], model="multiplicative")
print('sum of residuals {:2.2f}'.format(result.resid.sum()))
ax = result.plot()
ax.set_figwidth(16)
ax.set_figheight(8)

# __Note:__ While residuals in additive model is around zero, in multiplicative model they are around one.

# As we can see, multiplicative approach looks much better. There is still a bit of pattern in residuals but it has improved significantly compared to additive approach.

# ## Exercise
#

#
# let's try these on real data, which will be much harder. Use the data below and perform the following analyses:
#
# - STL decomposition using additive method.
#
#
# Hints
#
# * Find the relevent code, copy it, and change it to use `df['target']`
#
#

#

df = block0 = pd.read_csv("../../data/processed/smartmeter/block_0.csv", parse_dates=['day'], index_col=['day'])[['energy_sum']]
df = df.groupby('day').mean().iloc[:100]
df = df.rename(columns={'energy_sum':'target'})
df.plot()
df

# Exercise 1
from statsmodels.tsa.seasonal import seasonal_decompose
result = seasonal_decompose(df['target'],model = 'add')
result.plot();
print('sum of residuals {:2.2f}'.format(result.resid.sum()))

#
#

#
# ### Solution
# See solution
#
#
#
# ```Python
# # Exercise 1
# from statsmodels.tsa.seasonal import seasonal_decompose
# result = seasonal_decompose(df['target'],model = 'add')
# result.plot();
# print(f'sum of residuals {result.resid.sum():2.2f}')
# ```
#
#
#

# In the next notebook we will look at how to do predictions, including seasonality

# ## Further Reading:
# - [Pandas time-series documentation](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html)
# - [Statsmodel documentation](https://www.statsmodels.org/)
# - [Forecasting: Principles and Practice](https://otexts.com/fpp2/)
# - [Exponential Smoothing](https://en.wikipedia.org/wiki/Exponential_smoothing)
# - [Gentle Introduction to Exponential smoothing](https://machinelearningmastery.com/exponential-smoothing-for-time-series-forecasting-in-python/)
#
#

import numpy,import statsmodels,from statsmodels


# Copyright (c) 2021 Microsoft Corporation. Licensed under the MIT license.
# Reimplemtened by Jianwei Yang (jianwyan@microsoft.com)
# Adapted from https://github.com/jwyang/graph-rcnn.pytorch (Jianwei Yang) and https://github.com/yikang-li/MSDN (Yikang Li)
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn import Parameter
import pdb

class Message_Passing_Unit_v2(nn.Module):
    def __init__(self, fea_size, filter_size=128):
        super(Message_Passing_Unit_v2, self).__init__()
        self.w = nn.Linear(fea_size, filter_size, bias=True)
        self.fea_size = fea_size
        self.filter_size = filter_size

    def forward(self, unary_term, pair_term):

        if unary_term.size()[0] == 1 and pair_term.size()[0] > 1:
            unary_term = unary_term.expand(pair_term.size()[0],
                                           unary_term.size()[1])
        if unary_term.size()[0] > 1 and pair_term.size()[0] == 1:
            pair_term = pair_term.expand(unary_term.size()[0],
                                         pair_term.size()[1])

        # print '[unary_term, pair_term]', [unary_term, pair_term]
        gate = self.w(F.relu(unary_term)) * self.w(F.relu(pair_term))
        gate = torch.sigmoid(gate.sum(1))
        # print 'gate', gate
        output = pair_term * gate.expand(gate.size()[0], pair_term.size()[1])

        return output

class Message_Passing_Unit_v1(nn.Module):
    def __init__(self, fea_size, filter_size=128):
        super(Message_Passing_Unit_v1, self).__init__()
        self.w = nn.Linear(fea_size * 2, filter_size, bias=True)
        self.fea_size = fea_size
        self.filter_size = filter_size

    def forward(self, unary_term, pair_term):

        if unary_term.size()[0] == 1 and pair_term.size()[0] > 1:
            unary_term = unary_term.expand(pair_term.size()[0],
                                           unary_term.size()[1])
        if unary_term.size()[0] > 1 and pair_term.size()[0] == 1:
            pair_term = pair_term.expand(unary_term.size()[0],
                                         pair_term.size()[1])

        # print '[unary_term, pair_term]', [unary_term, pair_term]
        gate = torch.cat([unary_term, pair_term], 1)
        gate = F.relu(gate)
        gate = torch.sigmoid(self.w(gate)).mean(1)
        # print 'gate', gate
        output = pair_term * gate.view(-1, 1).expand(gate.size()[0],
                                                     pair_term.size()[1])

        return output

class Gated_Recurrent_Unit(nn.Module):
    def __init__(self, fea_size, dropout):
        super(Gated_Recurrent_Unit, self).__init__()
        self.wih = nn.Linear(fea_size, fea_size, bias=True)
        self.whh = nn.Linear(fea_size, fea_size, bias=True)
        self.dropout = dropout

    def forward(self, input, hidden):
        output = self.wih(F.relu(input)) + self.whh(F.relu(hidden))
        if self.dropout:
            output = F.dropout(output, training=self.training)
        return output

class MSDN_BASE(nn.Module):
    def __init__(self, fea_size, dropout=False, gate_width=128,
                 use_region=False, use_kernel_function=False,
                 update_step=2):
        super(MSDN_BASE, self).__init__()
        if update_step > 0:
            if use_kernel_function:
                Message_Passing_Unit = Message_Passing_Unit_v2
            else:
                Message_Passing_Unit = Message_Passing_Unit_v1

            self.gate_sub2pred = Message_Passing_Unit(fea_size, gate_width)
            self.gate_obj2pred = Message_Passing_Unit(fea_size, gate_width)
            self.gate_pred2sub = Message_Passing_Unit(fea_size, gate_width)
            self.gate_pred2obj = Message_Passing_Unit(fea_size, gate_width)

            self.GRU_object = Gated_Recurrent_Unit(fea_size,
                                                   dropout)  # nn.GRUCell(fea_size, fea_size) #
            self.GRU_pred = Gated_Recurrent_Unit(fea_size, dropout)

    def forward(self, feature_obj, feature_phrase, feature_region, mps_object,
                mps_phrase, mps_region):
        raise Exception('Please implement the forward function')

    # Here, we do all the operations outof loop, the loop is just to combine the features
    # Less kernel evoke frequency improve the speed of the model
    def prepare_message(self, target_features, source_features, select_mat,
                        gate_module):
        feature_data = []

        # transfer_list = np.where(select_mat > 0)

        if select_mat.data.sum() == 0:
            temp = Variable(torch.zeros(target_features.size()[1:]),
                            requires_grad=True).type_as(target_features)
            feature_data.append(temp)
        else:
            transfer_list = (select_mat.data > 0).nonzero(as_tuple=False)
            source_indices = Variable(transfer_list[:, 1])
            target_indices = Variable(transfer_list[:, 0])
            source_f = torch.index_select(source_features, 0, source_indices)
            target_f = torch.index_select(target_features, 0, target_indices)
            transferred_features = gate_module(target_f, source_f)

            for f_id in range(target_features.size()[0]):
                if select_mat[f_id, :].data.sum() > 0:
                    feature_indices = (transfer_list[:, 0] == f_id).nonzero(as_tuple=False)[0]
                    indices = Variable(feature_indices)
                    features = torch.index_select(transferred_features, 0,
                                                  indices).mean(0).view(-1)
                    feature_data.append(features)
                else:
                    temp = Variable(torch.zeros(target_features.size()[1:]),
                                    requires_grad=True).type_as(target_features)
                    feature_data.append(temp)

        return torch.stack(feature_data, 0)


# Parameters:
# → image_base64 (encoding image array)
# → image_url (s3 download authorization)

import os
import logging
import requests
import random

import tensorflow as tf
import tensorflow.keras.preprocessing.image as tfi
import pandas as pd
import numpy as np

from packages.lambda_handlers.response.wrapper import *
from packages.lambda_handlers.errors import *
from packages.lambda_handlers.handlers.lambda_handler import LambdaHandler

# Context:
# - Add EFS Path to python
# - Import libraries
# - Load model

# Handler
# - Read image url
# - Prepare image -> data augmentation
# - Make inference
# - Return predicted result

logger = logging.getLogger()
logger.setLevel(logging.INFO)

efs_path = "./"

"""
model_path = os.path.join(efs_path, 'model/')
model = tf.keras.models.load_model(model_path)
"""

#model_path = r"model.h5"
path = "/"
model_path = os.path.join(path, 'model/')
model = tf.keras.models.load_model(model_path)

logger.info("Lambda context execution success!")
print("Lambda context execution success!")

def lambda_handler(event, context):
    lambdaHandler = LambdaHandler(event=event, context=context)

    route_key = event['routeKey']
    image_url = event["queryStringParameters"]["image_url"]

    try:
        img_array = prepare_image(image_url, (128,128))

        predict = model.predict(img_array)

        prediction = random.randint(0,1)
        accuracy = 0.78945
        label = get_label(prediction)

        return buildResponse(operation=route_key,
                            data={  # Data predicted
                                "label": label,
                                "prediction": prediction,
                                "accuracy": accuracy
                            },
                            lambdaHandler=lambdaHandler).asjson()
    except:
        lambdaHandler.performError(BadRequestError())
        return buildResponse(operation=route_key,
                             data={},
                             lambdaHandler=lambdaHandler
                             ).asjson()

def prepare_image(url, IMG_SHAPE):
    # Preprocessing image from url -> tf.keras.preprocessing.image

    image_url = tf.keras.utils.get_file('Court', origin=url)
    img = tf.keras.preprocessing.image.load_img(image_url, target_size=IMG_SHAPE)
    os.remove(image_url) # Remove the cached file
    img_array = tf.keras.preprocessing.image.img_to_array(img)

    return img_array

def get_label(prediction):
    lbl = 'Cat' if prediction == 0 else 'Dog'
    return "It's a {lbl}".format(lbl=lbl)

def download_image(url):
    try:
        request = requests.get(url)
        print("Downloading image from url: {}".format(url))
        return request
    except requests.exceptions.ConnectionError as e:
        return e


# -*- coding: utf-8 -*-

__author__ = "Ngoc Huynh Bao"
__email__ = "ngoc.huynh.bao@nmbu.no"

import os
import h5py
import numpy as np
import warnings

from deoxys.keras.callbacks import CSVLogger
from ..model.callbacks import DeoxysModelCheckpoint, PredictionCheckpoint, \
    DBLogger
from ..model import model_from_full_config, model_from_config, load_model
from deoxys_vis import plot_log_performance_from_csv, mask_prediction, \
    plot_images_w_predictions, read_csv
from ..database import Tables, ExperimentAttr, HDF5Attr, SessionAttr, \
    SessionStatus

class Experiment:
    MODEL_PATH = '/model'
    MODEL_NAME = '/model.{epoch:03d}.h5'
    BEST_MODEL_PATH = '/best'
    PREDICTION_PATH = '/prediction'
    PREDICTION_NAME = '/prediction.{epoch:03d}.h5'
    LOG_FILE = '/logs.csv'
    PERFORMANCE_PATH = '/performance'
    PREDICTED_IMAGE_PATH = '/images'
    TEST_OUTPUT_PATH = '/test'
    PREDICT_TEST_NAME = '/prediction_test.h5'
    _max_size = 1

    def __init__(self,
                 log_base_path='logs',
                 best_model_monitors='val_loss',
                 best_model_modes='auto'):

        self.model = None

        self.architecture = None
        self.input_params = None
        self.model_params = None
        self.train_params = None
        self.data_reader = None
        self.weights_file = None

        self.log_base_path = log_base_path

        self.best_model_monitors = best_model_monitors if type(
            best_model_monitors) == list else [best_model_monitors]
        self.best_model_modes = best_model_modes if type(
            best_model_modes) == list else \
            [best_model_modes] * len(best_model_monitors)

        for i, (monitor, mode) in enumerate(zip(self.best_model_monitors,
                                                self.best_model_modes)):
            if mode not in ['auto', 'min', 'max']:
                warnings.warn('ModelCheckpoint mode %s is unknown, '
                              'fallback to auto mode.' % (mode),
                              RuntimeWarning)
                mode = 'auto'

            if mode == 'auto':
                if 'acc' in monitor or \
                        monitor.startswith('fmeasure') or \
                        'fbeta' in monitor:
                    self.best_model_modes[i] = 'max'
                else:
                    self.best_model_modes[i] = 'min'

    def from_full_config(self, file, weights_file=None, **kwargs):
        self.model = model_from_full_config(file,
                                            weights_file=weights_file,
                                            **kwargs)

        return self

    def from_config(self, architecture, input_params,
                    model_params, train_params,
                    dataset_params, weights_file=None):
        self.model = model_from_config(architecture, input_params,
                                       model_params, train_params,
                                       dataset_params,
                                       weights_file=weights_file)

        return self

    def from_file(self, filename):
        self.model = load_model(filename)

        return self

    def best_model(self):
        res = {}
        logger_path = self.log_base_path + self.LOG_FILE
        if os.path.isfile(logger_path):
            df = read_csv(logger_path, index_col='epoch',
                          usecols=['epoch'] + self.best_model_monitors)
            min_df = df.min()
            min_epoch = df.idxmin()
            max_df = df.max()
            max_epoch = df.idxmax()
            for monitor, mode in zip(self.best_model_monitors,
                                     self.best_model_modes):
                if mode == 'min':
                    val = min_df[monitor]
                    epoch = min_epoch[monitor]
                else:
                    val = max_df[monitor]
                    epoch = max_epoch[monitor]
                res[monitor] = {
                    'best': {
                        'val': val,
                        'epoch': epoch + 1
                    }}
        else:
            warnings.warn('No log files to check for best model')

        return res

    def run_experiment(self, train_history_log=True,
                       model_checkpoint_period=0,
                       prediction_checkpoint_period=0,
                       save_origin_images=False,
                       verbose=1,
                       epochs=None, initial_epoch=None, **custom_kwargs
                       ):
        log_base_path = self.log_base_path

        if self._check_run():
            if not os.path.exists(log_base_path):
                os.makedirs(log_base_path)

            kwargs = custom_kwargs or {}

            csv_logger_append = False

            if epochs:
                kwargs['epochs'] = epochs
            if initial_epoch:
                kwargs['initial_epoch'] = initial_epoch
                if initial_epoch > 0:
                    csv_logger_append = True

            callbacks = []

            if train_history_log:
                callback = self._create_logger(log_base_path,
                                               append=csv_logger_append)
                callbacks.append(callback)

            if model_checkpoint_period > 0:
                if not os.path.exists(log_base_path + self.MODEL_PATH):
                    os.makedirs(log_base_path + self.MODEL_PATH)

                callback = self._create_model_checkpoint(
                    log_base_path,
                    period=model_checkpoint_period)
                callbacks.append(callback)

            if prediction_checkpoint_period > 0:
                if not os.path.exists(log_base_path + self.PREDICTION_PATH):
                    os.makedirs(log_base_path + self.PREDICTION_PATH)

                callback = self._create_prediction_checkpoint(
                    log_base_path,
                    prediction_checkpoint_period,
                    use_original=save_origin_images
                )
                callbacks.append(callback)

            kwargs['callbacks'] = callbacks

            self.model.fit_train(**kwargs)

            return self

    def _plot_performance(self, log_base_path):
        if not os.path.exists(log_base_path + self.PERFORMANCE_PATH):
            os.makedirs(log_base_path + self.PERFORMANCE_PATH)

        if os.path.exists(log_base_path + self.LOG_FILE):
            print('\nPlotting performance metrics...')

            plot_log_performance_from_csv(
                filepath=log_base_path + self.LOG_FILE,
                output_path=log_base_path + self.PERFORMANCE_PATH)
        else:
            raise Warning('No log files for plotting performance')

    def plot_performance(self):
        log_base_path = self.log_base_path

        self._plot_performance(log_base_path)

        return self

    def _plot_prediction(self, log_base_path,
                         masked_images,
                         contour=True,
                         base_image_name='x',
                         truth_image_name='y',
                         predicted_image_name='predicted',
                         predicted_image_title_name='Image {index:05d}',
                         img_name='{index:05d}.png'):
        if os.path.exists(log_base_path + self.PREDICTION_PATH):
            print('\nCreating prediction images...')
            # mask images
            prediced_image_path = log_base_path + self.PREDICTED_IMAGE_PATH
            if not os.path.exists(prediced_image_path):
                os.makedirs(prediced_image_path)

            for filename in os.listdir(
                    log_base_path + self.PREDICTION_PATH):
                if filename.endswith(".h5") or filename.endswith(".hdf5"):
                    # Create a folder for storing result in that period
                    images_path = prediced_image_path + '/' + filename
                    if not os.path.exists(images_path):
                        os.makedirs(images_path)

                    self._plot_predicted_images(
                        data_path=log_base_path + self.PREDICTION_PATH
                        + '/' + filename,
                        out_path=images_path,
                        images=masked_images,
                        base_image_name=base_image_name,
                        truth_image_name=truth_image_name,
                        predicted_image_name=predicted_image_name,
                        title=predicted_image_title_name,
                        contour=contour,
                        name=img_name)

    def plot_prediction(self, masked_images,
                        contour=True,
                        base_image_name='x',
                        truth_image_name='y',
                        predicted_image_name='predicted',
                        predicted_image_title_name='Image {index:05d}',
                        img_name='{index:05d}.png'):
        log_base_path = self.log_base_path

        self._plot_prediction(
            log_base_path,
            masked_images,
            contour=contour,
            base_image_name=base_image_name,
            truth_image_name=truth_image_name,
            predicted_image_name=predicted_image_name,
            predicted_image_title_name=predicted_image_title_name,
            img_name=img_name)

        return self

    def _predict_test(self, filepath, use_original_image=False):
        data_info = self.model.data_reader.test_generator.description
        total_size = np.product(
            data_info[0]['shape']) * data_info[0]['total'] / 1e9

        # predict directly for data of size < max_size (1GB)
        if len(data_info) == 1 and total_size < self._max_size:
            predicted = self.model.predict_test(verbose=1)
            # Create the h5 file
            hf = h5py.File(filepath, 'w')
            hf.create_dataset('predicted', data=predicted)
            hf.close()

            if use_original_image:
                original_data = self.model.data_reader.original_test

                for key, val in original_data.items():
                    hf = h5py.File(filepath, 'a')
                    hf.create_dataset(key, data=val)
                    hf.close()
            else:
                # Create data from test_generator
                x = None
                y = None

                test_gen = self.model.data_reader.test_generator
                data_gen = test_gen.generate()

                for _ in range(test_gen.total_batch):
                    next_x, next_y = next(data_gen)
                    if x is None:
                        x = next_x
                        y = next_y
                    else:
                        x = np.concatenate((x, next_x))
                        y = np.concatenate((y, next_y))

                hf = h5py.File(filepath, 'a')
                hf.create_dataset('x', data=x)
                hf.create_dataset('y', data=y)
                hf.close()

        # for large data of same size, predict each chunk
        elif len(data_info) == 1:
            test_gen = self.model.data_reader.test_generator
            data_gen = test_gen.generate()

            next_x, next_y = next(data_gen)
            predicted = self.model.predict(next_x, verbose=1)

            input_shape = (data_info[0]['total'],) + data_info[0]['shape']
            input_chunks = (1,) + data_info[0]['shape']
            target_shape = (data_info[0]['total'],) + next_y.shape[1:]
            target_chunks = (1,) + next_y.shape[1:]

            with h5py.File(filepath, 'w') as hf:
                hf.create_dataset('x',
                                  shape=input_shape, chunks=input_chunks,
                                  compression='gzip')
                hf.create_dataset('y',
                                  shape=target_shape, chunks=target_chunks,
                                  compression='gzip')

                hf.create_dataset('predicted',
                                  shape=target_shape, chunks=target_chunks,
                                  compression='gzip')

            with h5py.File(filepath, 'a') as hf:
                next_index = len(next_x)
                hf['x'][:next_index] = next_x
                hf['y'][:next_index] = next_y
                hf['predicted'][:next_index] = predicted

            for _ in range(test_gen.total_batch - 1):
                next_x, next_y = next(data_gen)
                predicted = self.model.predict(next_x, verbose=1)

                curr_index = next_index
                next_index = curr_index + len(next_x)

                with h5py.File(filepath, 'a') as hf:
                    hf['x'][curr_index:next_index] = next_x
                    hf['y'][curr_index:next_index] = next_y
                    hf['predicted'][curr_index:next_index] = predicted

        # data of different size
        else:
            test_gen = self.model.data_reader.test_generator
            data_gen = test_gen.generate()

            for curr_info_idx, info in enumerate(data_info):
                next_x, next_y = next(data_gen)
                predicted = self.model.predict(next_x, verbose=1)

                input_shape = (info['total'],) + info['shape']
                input_chunks = (1,) + info['shape']
                target_shape = (info['total'],) + next_y.shape[1:]
                target_chunks = (1,) + next_y.shape[1:]
                if curr_info_idx == 0:
                    mode = 'w'
                else:
                    mode = 'a'
                with h5py.File(filepath, mode) as hf:
                    hf.create_dataset(f'{curr_info_idx:02d}/x',
                                      shape=input_shape,
                                      chunks=input_chunks,
                                      compression='gzip')
                    hf.create_dataset(f'{curr_info_idx:02d}/y',
                                      shape=target_shape,
                                      chunks=target_chunks,
                                      compression='gzip')

                    hf.create_dataset(f'{curr_info_idx:02d}/predicted',
                                      shape=target_shape,
                                      chunks=target_chunks,
                                      compression='gzip')

                with h5py.File(filepath, 'a') as hf:
                    next_index = len(next_x)
                    hf[f'{curr_info_idx:02d}/x'][:next_index] = next_x
                    hf[f'{curr_info_idx:02d}/y'][:next_index] = next_y
                    hf[f'{curr_info_idx:02d}/predicted'][
                        :next_index] = predicted

                while next_index < info['total']:
                    next_x, next_y = next(data_gen)
                    predicted = self.model.predict(next_x, verbose=1)

                    curr_index = next_index
                    next_index = curr_index + len(next_x)

                    with h5py.File(filepath, 'a') as hf:
                        hf[f'{curr_info_idx:02d}/x'][
                            curr_index:next_index] = next_x
                        hf[f'{curr_info_idx:02d}/y'][
                            curr_index:next_index] = next_y
                        hf[f'{curr_info_idx:02d}/predicted'][
                            curr_index:next_index] = predicted

    def run_test(self, use_best_model=False,
                 masked_images=None,
                 use_original_image=False,
                 contour=True,
                 base_image_name='x',
                 truth_image_name='y',
                 predicted_image_name='predicted',
                 image_name='{index:05d}.png',
                 image_title_name='Image {index:05d}'):

        log_base_path = self.log_base_path

        test_path = log_base_path + self.TEST_OUTPUT_PATH

        if not os.path.exists(test_path):
            os.makedirs(test_path)

        if use_best_model:
            raise NotImplementedError
        else:
            score = self.model.evaluate_test(verbose=1)
            print(score)
            filepath = test_path + self.PREDICT_TEST_NAME

            self._predict_test(filepath, use_original_image=use_original_image)

            if masked_images:
                self._plot_predicted_images(
                    data_path=filepath,
                    out_path=test_path,
                    images=masked_images,
                    base_image_name=base_image_name,
                    truth_image_name=truth_image_name,
                    predicted_image_name=predicted_image_name,
                    title=image_title_name,
                    contour=contour,
                    name=image_name)

        return self

    def run_lambda(self, lambda_fn, **kwargs):
        """
        Custom action between experiments
        """
        lambda_fn(self, **kwargs)
        return self

    def _create_logger(self, base_path, append=False):
        return CSVLogger(filename=base_path + self.LOG_FILE, append=append)

    def _create_model_checkpoint(self, base_path, period):
        return DeoxysModelCheckpoint(
            period=period,
            filepath=base_path + self.MODEL_PATH + self.MODEL_NAME)

    def _create_prediction_checkpoint(self, base_path, period, use_original):
        return PredictionCheckpoint(
            filepath=base_path + self.PREDICTION_PATH + self.PREDICTION_NAME,
            period=period, use_original=use_original)

    def _plot_predicted_images(self, data_path, out_path, images,
                               contour=True,
                               base_image_name='x',
                               truth_image_name='y',
                               predicted_image_name='predicted',
                               title='Image {index:05d}',
                               name='{index:05d}.png'):
        hf = h5py.File(data_path, 'r')
        keys = list(hf.keys())
        while 'create_group' in dir(hf[keys[0]]):
            new_keys = []
            for key in keys:
                new_keys.extend([f'{key}/{k}' for k in hf[key].keys()])
            keys = new_keys

        for index in images:
            kwargs = {key: hf[key][index] for key in keys}
            if base_image_name not in kwargs:
                for key in kwargs:
                    if key.endswith(base_image_name):
                        break
                prefix = key[:-len(base_image_name)]
                base_image_name = prefix + base_image_name
                truth_image_name = prefix + truth_image_name
                predicted_image_name = prefix + predicted_image_name

            img_name = out_path + '/' + name.format(index=index, **kwargs)
            try:
                if contour:
                    mask_prediction(img_name,
                                    image=kwargs[base_image_name],
                                    true_mask=kwargs[truth_image_name],
                                    pred_mask=kwargs[predicted_image_name],
                                    title=title.format(index=index, **kwargs))
                else:
                    plot_images_w_predictions(
                        img_name,
                        image=kwargs[base_image_name],
                        true_mask=kwargs[truth_image_name],
                        pred_mask=kwargs[predicted_image_name],
                        title=title.format(index=index, **kwargs))
            except Exception as e:
                print('An error occurred while plotting prediction', e)

    def _check_run(self):
        if self.model:
            if self.model._data_reader:
                if self.model.is_compiled:
                    return True
        raise RuntimeError("Cannot run experiment with incomplete model")
        return False

class ExperimentDB(Experiment):  # pragma: no cover

    def __init__(self, dbclient, experiment_id=None, session_id=None,
                 log_base_path='logs',
                 best_model_monitors='val_loss',
                 best_model_modes='auto'):
        """
        An experiment logging performance to a database

        Parameters
        ----------
        dbclient : deoxys.database.DBClient
            The database client
        experiment_id : str, int, or ObjectID the dbclient, optional
            Experiment id, by default None
        session_id : str, int, or ObjectID depending of the dbclient, optional
            Session id, by default None
        log_base_path : str, optional
            Base path to log files, by default 'logs'
        best_model_monitors : str, optional
            Attribute to monitor, by default 'val_loss'
        best_model_modes : str, optional
            One of 'max', 'min', 'auto', by default 'auto'

        Raises
        ------
        ValueError
            When both `session_id` and `experiment_id` are not set
        """
        super().__init__(log_base_path, best_model_monitors, best_model_modes)

        self.dbclient = dbclient

        if experiment_id is None and session_id is None:
            raise ValueError('"session_id" or "experiment_id" must be set')

        if session_id:
            last_model = self.dbclient.find_max(Tables.MODELS, {
                HDF5Attr.SESSION_ID: session_id}, HDF5Attr.EPOCH)
            self.curr_epoch = last_model[HDF5Attr.EPOCH]
            self.session = dbclient.find_by_id(Tables.SESSIONS, session_id)

            self._update_epoch(last_model[HDF5Attr.EPOCH])

            self.from_file(last_model[HDF5Attr.FILE_LOCATION])
        else:
            insert_res = dbclient.insert(Tables.SESSIONS, {
                SessionAttr.EXPERIMENT_ID: experiment_id,
                SessionAttr.CURRENT_EPOCH: 0,
                SessionAttr.STATUS: 'created'
            }, time_logs=True)

            self.session = self.dbclient.find_by_id(
                Tables.SESSIONS, insert_res.inserted_id)

            self.curr_epoch = 0

            experiment_obj = self.dbclient.find_by_id(
                Tables.EXPERIMENTS, experiment_id)

            if ExperimentAttr.CONFIG in experiment_obj:
                self.from_full_config(experiment_obj[ExperimentAttr.CONFIG])
            elif ExperimentAttr.SAVED_MODEL_LOC in experiment_obj:
                self.from_file(experiment_obj[ExperimentAttr.SAVED_MODEL_LOC])

        self.log_base_path = os.path.join(
            log_base_path, str(self.dbclient.get_id(self.session)))

    def run_experiment(self, train_history_log=True,
                       model_checkpoint_period=0,
                       prediction_checkpoint_period=0,
                       save_origin_images=False,
                       verbose=1,
                       epochs=None
                       ):
        log_base_path = self.log_base_path

        if self._check_run():
            if not os.path.exists(log_base_path):
                os.makedirs(log_base_path)

            kwargs = {}

            if epochs:
                kwargs['epochs'] = epochs + self.curr_epoch

            kwargs['initial_epoch'] = self.curr_epoch

            callbacks = []

            if train_history_log:
                callback = self._create_logger(log_base_path,
                                               append=self.curr_epoch > 0)
                callbacks.append(callback)

                callback = self._create_db_logger()
                callbacks.append(callback)

            if model_checkpoint_period > 0:
                if not os.path.exists(log_base_path + self.MODEL_PATH):
                    os.makedirs(log_base_path + self.MODEL_PATH)

                callback = self._create_model_checkpoint(
                    log_base_path,
                    period=model_checkpoint_period)
                callbacks.append(callback)

            if prediction_checkpoint_period > 0:
                if not os.path.exists(log_base_path + self.PREDICTION_PATH):
                    os.makedirs(log_base_path + self.PREDICTION_PATH)

                callback = self._create_prediction_checkpoint(
                    log_base_path,
                    prediction_checkpoint_period,
                    use_original=save_origin_images
                )
                callbacks.append(callback)

            kwargs['callbacks'] = callbacks

            self._update_status(SessionStatus.TRAINING)

            try:
                self.model.fit_train(**kwargs)

                self.curr_epoch += epochs
                self._update_status(SessionStatus.FINISHED)
                self._update_epoch(self.curr_epoch)
            except Exception:
                self._update_status(SessionStatus.FAILED)

            return self

    def _create_db_logger(self):
        return DBLogger(self.dbclient, self.dbclient.get_id(self.session))

    def _update_epoch(self, new_epoch):
        self.dbclient.update_by_id(
            Tables.SESSIONS, self.dbclient.get_id(self.session),
            {SessionAttr.CURRENT_EPOCH: new_epoch}, time_logs=True)

    def _update_status(self, new_status):
        self.dbclient.update_by_id(
            Tables.SESSIONS, self.dbclient.get_id(self.session),
            {SessionAttr.STATUS: new_status}, time_logs=True)

    def _create_model_checkpoint(self, base_path, period):
        return DeoxysModelCheckpoint(
            period=period,
            filepath=base_path + self.MODEL_PATH + self.MODEL_NAME,
            dbclient=self.dbclient,
            session=self.dbclient.get_id(self.session))

    def _create_prediction_checkpoint(self, base_path, period, use_original):
        return PredictionCheckpoint(
            filepath=base_path + self.PREDICTION_PATH + self.PREDICTION_NAME,
            period=period, use_original=use_original,
            dbclient=self.dbclient,
            session=self.dbclient.get_id(self.session))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import torch
import torch.nn.parallel
from torch.autograd import Variable
import torch.optim as optim
import random
import sys
import os
import numpy as np
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
sys.path.append('..')
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, 'models')))
sys.path.append(os.path.abspath(os.path.join(BASE_DIR, 'dataloaders')))
import shapenet_part_loader
from utils import PointLoss
from utils import distance_squre
from pointcapsnet_ae import PointCapsNet
import data_utils as d_utils

parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=1, help='input batch size')
parser.add_argument('--n_epochs', type=int, default=300, help='number of epochs to train for')
parser.add_argument('--workers', type=int,default=2, help='number of data loading workers')
parser.add_argument('--prim_caps_size', type=int, default=1024, help='number of primary point caps')
parser.add_argument('--prim_vec_size', type=int, default=16, help='scale of primary point caps')
parser.add_argument('--latent_caps_size', type=int, default=64, help='number of latent caps')
parser.add_argument('--latent_vec_size', type=int, default=64, help='scale of latent caps')
parser.add_argument('--num_points', type=int, default=2048, help='input point set size')
parser.add_argument('--crop_point_num',type=int,default=512,help='number of crop points ')
parser.add_argument('--outf', type=str, default='tmp_checkpoints', help='output folder')
parser.add_argument('--model', type=str, default='Trained_Recon_Model_Capsule/3dCapsule120.pth', help='model path')
parser.add_argument('--dataset', type=str, default='shapenet_part', help='dataset: shapenet_part, shapenet_core13, shapenet_core55')
opt = parser.parse_args()
print(opt)

test_dset = shapenet_part_loader.PartDataset( root='../dataset/shapenetcore_partanno_segmentation_benchmark_v0/',classification=True, class_choice='Motorbike', npoints=opt.num_points, split='test')
test_dataloader = torch.utils.data.DataLoader(test_dset, batch_size=opt.batch_size,
                                         shuffle=True,num_workers = int(opt.workers))

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
capsule_net = PointCapsNet(opt.prim_caps_size, opt.prim_vec_size, opt.latent_caps_size, opt.latent_caps_size, opt.num_points)
capsule_net.load_state_dict(torch.load(opt.model,map_location=lambda storage, location: storage)['state_dict'])
print("Let's use", torch.cuda.device_count(), "GPUs!")
capsule_net.to(device)
capsule_net = torch.nn.DataParallel(capsule_net)
capsule_net.eval()

criterion_PointLoss = PointLoss().to(device)
errG_min = 100
n = 0
for i, data in enumerate(test_dataloader, 0):

    real_point, target = data
    real_point = torch.unsqueeze(real_point, 1)
    input_cropped =torch.FloatTensor(opt.batch_size, 1, opt.num_points-opt.crop_point_num, 3)
    input_cropped_ours =torch.FloatTensor(opt.batch_size, 1, opt.num_points, 3)
    input_cropped_ours.resize_(real_point.size()).copy_(real_point)
    batch_size = real_point.size()[0]
    p_origin = [0,0,0]
    choice =[torch.Tensor([1,0,0]),torch.Tensor([0,0,1]),torch.Tensor([1,0,1])]
#    points = [x for x in range(0,opt.pnum-1)]
#    choice =random.sample(points,5)
    index = random.sample(choice,1)
    distance_list = []
#    p_center  = real_point[0,0,index]
    p_center = index[0]
    for num in range(opt.num_points):
        distance_list.append(distance_squre(real_point[0,0,num],p_center))
    distance_order = sorted(enumerate(distance_list), key = lambda x:x[1])

    for sp in range(opt.crop_point_num):
        input_cropped_ours.data[0,0,distance_order[sp][0]] = torch.FloatTensor([0,0,0])

    crop_num_list = []
    for num in range(opt.num_points-opt.crop_point_num):
        crop_num_list.append(distance_order[num+opt.crop_point_num][0])
    indices = torch.LongTensor(crop_num_list)
    input_cropped[0,0]=torch.index_select(real_point[0,0],0,indices)

    real_point = torch.squeeze(real_point,1)
    input_cropped = torch.squeeze(input_cropped,1)
    input_cropped = input_cropped.transpose(2, 1)
    real_point = real_point.to(device)
    input_cropped = input_cropped.to(device)
    codewords,reconstruction=capsule_net(input_cropped)
    reconstruction_ = reconstruction.transpose(2, 1).contiguous()
    reconstruction_ = reconstruction_.cuda()
    input_cropped = input_cropped.transpose(2, 1)
    real_point  = real_point.cuda()
#    real_center_key =real_center_key.cuda()
    errG = criterion_PointLoss(reconstruction_,real_point)
    errG = errG.cpu()
    if errG.detach().numpy()>errG_min:
        pass

    else:
        errG_min = errG.detach().numpy()
        print(errG_min)
        reconstruction_ =reconstruction_.cpu()
        np_recon = reconstruction_.data[0].detach().numpy()
        input_cropped = input_cropped.cpu()
        np_crop = input_cropped.data[0].detach().numpy()
#            num_real = np.array([[len(np_real)]])
        real_point = real_point.cpu()
        np_real = real_point.data[0].detach().numpy()

        input_cropped_ours = torch.squeeze(input_cropped_ours,1)
        input_cropped_ours = input_cropped_ours.cpu()
        np_crop_ours = input_cropped_ours[0].detach().numpy()

        n=n+1
        k = 0
#            num_crop = np.array([[len(np_crop)]])
        #np_ini = real_point[0,0].detach().numpy() #1024
        #np_crop = np.
        np.savetxt('test_example_capsule/crop'+str(n)+'.csv', np_crop, fmt = "%f,%f,%f")
        np.savetxt('test_example_capsule/recon'+str(n)+'.csv', np_recon, fmt = "%f,%f,%f")
        np.savetxt('test_example_capsule/real'+str(n)+'.csv', np_real, fmt = "%f,%f,%f")
        np.savetxt('test_example_capsule/crop-ours'+str(n)+'.csv', np_crop_ours, fmt = "%f,%f,%f")


from __future__ import division

import os
import random
import sys
import logging

import cv2
import fire
import numpy as np
import tensorflow as tf
import tensorflow.contrib.metrics as metrics
import time
import pickle

import commons
from configs import VGGConf
from networks import VGG16
from pystopwatch import StopWatchManager

_log_level = logging.DEBUG
_logger = logging.getLogger('ADNetRunner')
_logger.setLevel(_log_level)
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(_log_level)
formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
ch.setFormatter(formatter)
_logger.addHandler(ch)

class VGGRunner:
    MAX_BATCHSIZE = 512

    def __init__(self):
        self.tensor_input = tf.placeholder(tf.float32, shape=(None, 128, 128, 3), name='input')
        self.tensor_lb = tf.placeholder(tf.int32, shape=(None, ), name='lb_action')
        self.tensor_is_training = tf.placeholder(tf.bool, name='is_training')
        self.learning_rate_placeholder = tf.placeholder(tf.float32, [], name='learning_rate')
        config = tf.ConfigProto() #log_device_placement=True)

        self.persistent_sess = tf.Session(config=config)

        self.model = VGG16(self.learning_rate_placeholder)
        self.model.create_network(self.tensor_input, self.tensor_lb, self.tensor_is_training)
        #self.persistent_sess.run(tf.global_variables_initializer())
        self.model.read_original_weights(self.persistent_sess)
        self.iteration = 0
        self.imgwh = None

        self.loss = []
        self.train, self.test =[], []
        self.train_lb, self.test_lb = [], []

        self.stopwatch = StopWatchManager()

    def read_train_test(self, data_path, is_train=True):
        path = commons.get_image_pathes(data_path)
        images, labels = commons.read_image_lb_path(path, is_train)
        train_f_num = VGGConf.g()['train_folder_num']
        test_f_num = VGGConf.g()['test_folder_num']
        f_num = train_f_num + test_f_num
        print(f_num)
        assert f_num == len(images)
        assert len(images) == len(labels)

        for idx in range(f_num):
            if idx < train_f_num:
                if not is_train: continue
                self.train.extend(images[idx])
                self.train_lb.extend(labels[idx])
                print("Train data is added")
            else:
                self.test.extend(images[idx])
                self.test_lb.extend(labels[idx])
                print("Test data is added")

        _temp = [0,0,0,0,0]
        for i, a in enumerate(self.train_lb):
            _temp[int(a)]+=1
        print("Trainind Data Label Distribution: ", _temp)
        print("READING DATA IS COMPLETED")

    def label_index(self):
        label_array = np.array(self.test_lb)
        self.lb_idx = []
        for i in range(5):
            self.lb_idx.append(list(np.where(label_array==i)))

    def run_train(self, data_path='./data'):
        train = pickle.load(open('train_img.txt','rb'))
        self.train, self.train_lb = [], []
        for lb in range(5):
            self.train.extend(train[lb])
            num_lb = len(train[lb])
            self.train_lb.extend([lb]*num_lb)
        #self.read_train_test(data_path)
        train_num = len(self.train)
        train_index = list(range(train_num))

        EPOCH = VGGConf.g()['epoch']
        BATCH_SIZE = VGGConf.g()['minibatch_size']
        BATCH_NUM = int(len(train_index)/BATCH_SIZE) - 1
        learning_rate = VGGConf.g()['learning_rate']
        print("START TO TRAIN")
        for epoch_iter in range(EPOCH):
            iter = 0
            #if epoch_iter != 0 and epoch_iter % 100 ==0 : learning_rate *= 0.5
            random.shuffle(train_index)
            #for i in range(5): random.shuffle(self.lb_idx[i])
            for batch_iter in range(BATCH_NUM):
                batch_img = commons.choices_by_idx(self.train, train_index[iter*BATCH_SIZE: (iter+1)*BATCH_SIZE])
                batch_lb = commons.choices_by_idx(self.train_lb, train_index[iter*BATCH_SIZE: (iter+1)*BATCH_SIZE])

                _, loss = self.persistent_sess.run(
                    [self.model.opt, self.model.loss],
                    feed_dict={
                        self.tensor_input: np.array(batch_img)/255.,
                        self.tensor_lb: batch_lb,
                        self.learning_rate_placeholder: learning_rate,
                        self.tensor_is_training: True
                    }
                )

                print("{}/{} BATCH LOSS: {}".format(batch_iter, BATCH_NUM, np.sum(loss)))
                iter +=1

        self.model.save_train_model(self.persistent_sess, VGGConf.g()['save_weight_dir'])
        self.run_test(weight_load=False)

    def run_test(self, weight_load=True, data_path='./data'):
        test = pickle.load(open('test_img.txt','rb'))
        self.test, self.test_lb = [], []
        for lb in range(5):
            self.test.extend(test[lb])
            num_lb = len(test[lb])
            self.test_lb.extend([lb]*num_lb)

        MAX_BATCHSIZE = self.MAX_BATCHSIZE
        if weight_load:
            #self.read_train_test(data_path, is_train=False)
            self.model.read_original_weights(self.persistent_sess)
        test_num = len(self.test)
        if test_num % self.MAX_BATCHSIZE == 0 :
            test_batch_num = int(test_num / self.MAX_BATCHSIZE)
        else:
            test_batch_num = int(test_num / self.MAX_BATCHSIZE) + 1

        test_output = []

        test_idx = list(range(test_num))
        test_iter = 0
        print("START TO TEST")
        for batch_iter in range(test_batch_num):
            if batch_iter == test_batch_num - 1:
                batch_img = commons.choices_by_idx(self.test, test_idx[test_iter*MAX_BATCHSIZE:])
                batch_lb = commons.choices_by_idx(self.test_lb, test_idx[test_iter*MAX_BATCHSIZE:])
            else:
                batch_img = commons.choices_by_idx(self.test, test_idx[test_iter*MAX_BATCHSIZE: (test_iter+1)*MAX_BATCHSIZE])
                batch_lb = commons.choices_by_idx(self.test_lb, test_idx[test_iter*MAX_BATCHSIZE: (test_iter+1)*MAX_BATCHSIZE])

            output = self.persistent_sess.run(
                self.model.output, feed_dict={
                        self.tensor_input: np.array(batch_img)/255.,
                        self.tensor_lb: batch_lb,
                        self.tensor_is_training: False
                }
            )
            test_output.extend(output)
            test_iter +=1

        pickle.dump(self.test_lb, open('./gt.txt', 'wb'))
        pickle.dump(test_output, open('./pred.txt', 'wb'))

        accuracy = 0.0
        for idx, output in enumerate(test_output):
            max_idx = np.argmax(output)
            if max_idx == int(self.test_lb[idx]): accuracy+=1.0

        accuracy = accuracy / (len(self.test_lb))
        print("Test Accuracy: ", accuracy)

    def _get_features(self, samples):
        feats = []
        for batch in commons.chunker(samples, ADNetRunner.MAX_BATCHSIZE):
            feats_batch = self.persistent_sess.run(self.adnet.layer_feat, feed_dict={
                self.adnet.input_tensor: batch
            })
            feats.extend(feats_batch)
        return feats

        # train_fc_finetune_hem
        self._finetune_fc(
            img, pos_boxes, neg_boxes, pos_lb_action,
            ADNetConf.get()['initial_finetune']['learning_rate'],
            ADNetConf.get()['initial_finetune']['iter']
        )

        self.histories.append((pos_boxes, neg_boxes, pos_lb_action, np.copy(img), self.iteration))
        _logger.info('ADNetRunner.initial_finetune t=%.3f' % t)
        self.stopwatch.stop('initial_finetune')

    def __del__(self):
        self.persistent_sess.close()

if __name__ == '__main__':
    VGGConf.get('./conf/conf.yaml')

    random.seed(1258)
    np.random.seed(1258)
    tf.set_random_seed(1258)

    fire.Fire(VGGRunner)


import os
import glob
import tqdm
import torch
import random
import librosa
import argparse
import numpy as np
from multiprocessing import Pool, cpu_count

from utils.audio import Audio
from utils.hparams import HParam

def formatter(dir_, form, num):
    return os.path.join(dir_, form.replace('*', '%06d' % num))

def vad_merge(w):
    intervals = librosa.effects.split(w, top_db=20)
    temp = list()
    for s, e in intervals:
        temp.append(w[s:e])
    return np.concatenate(temp, axis=None)

def mix(hp, args, num, s1_target, s2, train):
    srate = hp.audio.sample_rate
    dir_ = os.path.join(args.out_dir, 'train' if train else 'test')

    # d, _ = librosa.load(s1_dvec, sr=srate)
    w1, _ = librosa.load(s1_target, sr=srate)
    w2, _ = librosa.load(s2, sr=srate)
    # assert   len(w1.shape) == len(w2.shape) == 1, \
    #     'wav files must be mono, not stereo'

    # d, _ = librosa.effects.trim(d, top_db=20)
    w1, _ = librosa.effects.trim(w1, top_db=20)
    w2, _ = librosa.effects.trim(w2, top_db=20)

    # if reference for d-vector is too short, discard it
    # if d.shape[0] < 1.1 * hp.embedder.window * hp.audio.hop_length:
    #     return

    # # LibriSpeech dataset have many silent interval, so let's vad-merge them
    # # VoiceFilter paper didn't do that. To test SDR in same way, don't vad-merge.
    if args.vad == 1:
        w1, w2 = vad_merge(w1), vad_merge(w2)

    # # I think random segment length will be better, but let's follow the paper first
    # # fit audio to `hp.data.audio_len` seconds.
    # # if merged audio is shorter than `L`, discard it
    L = int(srate * hp.data.audio_len)
    if w1.shape[0] < L or w2.shape[0] < L:
        return
    w1, w2 = w1[:L], w2[:L]

    mixed = w1 + (w2*0.1)

    norm = np.max(np.abs(mixed)) * 1.1
    w1, w2, mixed = w1/norm, w2/norm, mixed/norm

    # save vad & normalized wav files
    # target_wav_path = formatter(dir_, hp.form.target.wav, num)
    mixed_wav_path = formatter(dir_, hp.form.mixed.wav, num)
    # librosa.output.write_wav(target_wav_path, w1, srate)
    librosa.output.write_wav(mixed_wav_path, mixed, srate)

    # save magnitude spectrograms
    # target_mag, _ = audio.wav2spec(w1)
    # mixed_mag, _ = audio.wav2spec(mixed)
    # target_mag_path = formatter(dir_, hp.form.target.mag, num)
    # mixed_mag_path = formatter(dir_, hp.form.mixed.mag, num)
    # torch.save(torch.from_numpy(target_mag), target_mag_path)
    # torch.save(torch.from_numpy(mixed_mag), mixed_mag_path)

    # save selected sample as text file. d-vec will be calculated soon
    # dvec_text_path = formatter(dir_, hp.form.dvec, num)
    # with open(dvec_text_path, 'w') as f:
    #     f.write(s1_dvec)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--config', type=str, required=True,
                        help="yaml file for configuration")
    parser.add_argument('-d', '--libri_dir', type=str, default=None,
                        help="Directory of LibriSpeech dataset, containing folders of train-clean-100, train-clean-360, dev-clean.")
    parser.add_argument('-v', '--voxceleb_dir', type=str, default=None,
                        help="Directory of VoxCeleb2 dataset, ends with 'aac'")
    parser.add_argument('-o', '--out_dir', type=str, required=True,
                        help="Directory of output training triplet")
    parser.add_argument('-p', '--process_num', type=int, default=None,
                        help='number of processes to run. default: cpu_count')
    parser.add_argument('--vad', type=int, default=0,
                        help='apply vad to wav file. yes(1) or no(0, default)')
    args = parser.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    os.makedirs(os.path.join(args.out_dir, 'train'), exist_ok=True)
    os.makedirs(os.path.join(args.out_dir, 'test'), exist_ok=True)

    hp = HParam(args.config)

    cpu_num = cpu_count() if args.process_num is None else args.process_num

    if args.libri_dir is None and args.voxceleb_dir is None:
        raise Exception("Please provide directory of data")

    if args.libri_dir is not None:
        train_folders = [x for x in glob.glob(os.path.join(args.libri_dir, 'train-clean-100', '*'))
                            if os.path.isdir(x)] + \
                        [x for x in glob.glob(os.path.join(args.libri_dir, 'train-clean-360', '*'))
                            if os.path.isdir(x)]
                        # we recommned to exclude train-other-500
                        # See https://github.com/mindslab-ai/voicefilter/issues/5#issuecomment-497746793
                        # + \
                        #[x for x in glob.glob(os.path.join(args.libri_dir, 'train-other-500', '*'))
                        #    if os.path.isdir(x)]
        test_folders = [x for x in glob.glob(os.path.join(args.libri_dir, 'dev-clean', '*'))]

    elif args.voxceleb_dir is not None:
        all_folders = [x for x in glob.glob(os.path.join(args.voxceleb_dir, '*'))
                            if os.path.isdir(x)]
        train_folders = all_folders[:-20]
        test_folders = all_folders[-20:]

    train_spk = [glob.glob(os.path.join(spk, '**', hp.form.input), recursive=True)
                    for spk in train_folders]
    train_spk = [x for x in train_spk if len(x) >= 2]

    test_spk = [glob.glob(os.path.join(spk, '**', hp.form.input), recursive=True)
                    for spk in test_folders]
    test_spk = [x for x in test_spk if len(x) >= 2]

    audio = Audio(hp)

    # def train_wrapper(num):
    #     spk1, spk2 = random.sample(train_spk, 2)
    #     s1_dvec, s1_target = random.sample(spk1, 2)
        # s2 = random.choice(spk2)
        # mix(hp, args, audio, num, s1_dvec, s1_target, s2, train=True)

    def test_wrapper(num):
        s1, s2 =random.sample(test_spk, 2)
        spk2 = random.choice(s2)
        for spk1s in  test_spk :
            for spk1 in spk1s:
                mix(hp, args, num, spk1, spk2, train=False)

    # arr = list(range(10**5))
    # with Pool(cpu_num) as p:
    #     r = list(tqdm.tqdm(p.imap(train_wrapper, arr), total=len(arr)))

    arr = list(range(10**2))
    with Pool(cpu_num) as p:
        r = list(tqdm.tqdm(p.imap(test_wrapper, arr), total=len(arr)))