# coding: utf-8

#
Table of Contents

#

1  Introduction

2  Setup

3  Read in the 5 km nearir and ir water vapor images

3.0.1  Here is wv_ir

3.0.2  Here is wv_nearir_ir mapped to the low resolution area_def

4  Questions

4.1  Question 9.1

4.2  Question 9.2

4.2.1  Here is the test that autograder will run

# # Introduction
#
# This assignment uses the functions and files defined in level2_cartopy_plot to look
# at the difference between wv_ir and wv_nearir mapped onto the low resolution
# grid.
#
# # Setup
#
# 1. Understand  and run level2_cartopy_plot
#
# 1. Setup should be the same, rerun test to check
#
# ```
#       python -m a301.install_tests.assign8_test
# ```
#

# In[1]:

from matplotlib import cm

import numpy as np
from matplotlib import pyplot as plt
from matplotlib.colors import Normalize
from IPython.display import Image,display
from pyresample import geometry
import pdb

#Image('figures/MYBRGB.A2016224.2100.006.2016237025650.jpg',width=600)

# In[2]:

get_ipython().run_line_magic('matplotlib', 'inline')
from matplotlib import cm
from a301.scripts.modismeta_read import parseMeta
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.colors import Normalize
from IPython.display import Image,display
import a301
from pathlib import Path
from pyhdf.SD import SD, SDC
import pprint
import json
from a301.geometry import area_def_from_dict, get_image, plot_image

# # Read in the 5 km nearir and ir water vapor images
#
#
# use get_image and plot_image from the library

# ### Here is wv_ir

# In[3]:

foldername=a301.map_dir / Path('wv_maps')
image_wv_ir, area_def_lr = get_image(foldername, 'wv_ir')

# In[4]:

fig,ax=plot_image(image_wv_ir, area_def_lr)
ax.set_title('5 km IR water vapor (cm)');

# ### Here is wv_nearir_ir mapped to the low resolution area_def

# In[5]:

image_wv_nearir_lr, area_def_hr = get_image(foldername, 'wv_nearir_lr')
fig,ax=plot_image(image_wv_nearir_lr, area_def_hr)
ax.set_title('1 km IR water vapor (cm) at 5k low resolution (lr)');
print(image_wv_nearir_lr.shape)

# # Questions

# ## Question 9.1
#
# in the cell below calculate the difference between the nearir_ir and wv_ir
# images as a new variable called wv_difference and plot
# it using plot_image
#
# Use a diverging palette like RdYlBu (see [matplotlib palettes](https://matplotlib.org/tutorials/colors/colormaps.html))
# to highlight the areas where the difference is zero and show postiive and negative departures.
# Set your vmin, vmax to -1, 1  cm.

# In[6]:

wv_difference = (image_wv_nearir_lr - image_wv_ir)
### BEGIN SOLUTION
pal = plt.get_cmap('RdYlBu').reversed()
fig,ax=plot_image(wv_difference, area_def_lr,vmin=-1,vmax=+1,palette=pal)
ax.set_title('1km - 5km wv (cm)');
### END SOLUTION

# ## Question 9.2
#
# Write out your wv_difference image and using dump_image

# In[7]:

from a301.geometry import dump_image, area_def_to_dict
image_name='wv_difference'
### BEGIN SOLUTION
metadata_dict=dict()
metadata_dict['area_def']=area_def_to_dict(area_def_lr)
metadata_dict['image_name']=image_name
metadata_dict['description']='column wv difference between nearir and ir retrievals'
metadata_dict['modis_metadata']=parseMeta(a301.data_dir / Path('myd05_l2_10_7.hdf'))
metadata_dict['history']='written by level2_cartopy_resample.ipynb'
out_folder = a301.map_dir / Path('wv_maps')
dump_image(wv_difference,metadata_dict,out_folder,image_name)
### END SOLUTION

# ### Here is the test that autograder will run
#
# Copy this code in a new cell and make sure it passes the test (i.e. doesn't complain)
# This is what autograder will run automatically

# In[9]:

from a301.geometry import get_image
#
# read back the image and the area_def
#
foldername= a301.map_dir / Path('wv_maps')
image_array_name='wv_difference'
image_array,area_def=get_image(foldername,image_array_name)
#
# make sure the area_def and image sizes agree
#
area_def_size=(area_def.y_size,area_def.x_size)
image_size=image_array.shape
assert(area_def_size == image_size)


#!/usr/bin/env python
# coding: utf-8

# # Residual Networks
#
# Welcome to the first assignment of this week! You'll be building a very deep convolutional network, using Residual Networks (ResNets). In theory, very deep networks can represent very complex functions; but in practice, they are hard to train. Residual Networks, introduced by [He et al.](https://arxiv.org/pdf/1512.03385.pdf), allow you to train much deeper networks than were previously feasible.
#
# **By the end of this assignment, you'll be able to:**
#
# - Implement the basic building blocks of ResNets in a deep neural network using Keras
# - Put together these building blocks to implement and train a state-of-the-art neural network for image classification
# - Implement a skip connection in your network
#
# For this assignment, you'll use Keras.
#
# Before jumping into the problem, run the cell below to load the required packages.

# ## Table of Content
#
# - [1 - Packages](#1)
# - [2 - The Problem of Very Deep Neural Networks](#2)
# - [3 - Building a Residual Network](#3)
#     - [3.1 - The Identity Block](#3-1)
#         - [Exercise 1 - identity_block](#ex-1)
#     - [3.2 - The Convolutional Block](#3-2)
#         - [Exercise 2 - convolutional_block](#ex-2)
# - [4 - Building Your First ResNet Model (50 layers)](#4)
#     - [Exercise 3 - ResNet50](#ex-3)
# - [5 - Test on Your Own Image (Optional/Ungraded)](#5)
# - [6 - Bibliography](#6)

#
# ## 1 - Packages

# In[1]:

import tensorflow as tf
import numpy as np
import scipy.misc
from tensorflow.keras.applications.resnet_v2 import ResNet50V2
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet_v2 import preprocess_input, decode_predictions
from tensorflow.keras import layers
from tensorflow.keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D, AveragePooling2D, MaxPooling2D, GlobalMaxPooling2D
from tensorflow.keras.models import Model, load_model
from resnets_utils import *
from tensorflow.keras.initializers import random_uniform, glorot_uniform, constant, identity
from tensorflow.python.framework.ops import EagerTensor
from matplotlib.pyplot import imshow

from test_utils import summary, comparator
import public_tests

get_ipython().run_line_magic('matplotlib', 'inline')

#
# ## 2 - The Problem of Very Deep Neural Networks
#
# Last week, you built your first convolutional neural networks: first manually with numpy, then using Tensorflow and Keras.
#
# In recent years, neural networks have become much deeper, with state-of-the-art networks evolving from having just a few layers (e.g., AlexNet) to over a hundred layers.
#
# * The main benefit of a very deep network is that it can represent very complex functions. It can also learn features at many different levels of abstraction, from edges (at the shallower layers, closer to the input) to very complex features (at the deeper layers, closer to the output).
#
# * However, using a deeper network doesn't always help. A huge barrier to training them is vanishing gradients: very deep networks often have a gradient signal that goes to zero quickly, thus making gradient descent prohibitively slow.
#
# * More specifically, during gradient descent, as you backpropagate from the final layer back to the first layer, you are multiplying by the weight matrix on each step, and thus the gradient can decrease exponentially quickly to zero (or, in rare cases, grow exponentially quickly and "explode," from gaining very large values).
#
# * During training, you might therefore see the magnitude (or norm) of the gradient for the shallower layers decrease to zero very rapidly as training proceeds, as shown below:

#
#
   Figure 1   : Vanishing gradient
 The speed of learning decreases very rapidly for the shallower layers as the network trains

#
# Not to worry! You are now going to solve this problem by building a Residual Network!

#
# ## 3 - Building a Residual Network
#
# In ResNets, a "shortcut" or a "skip connection" allows the model to skip layers:
#
#
#
   Figure 2   : A ResNet block showing a skip-connection

#
# The image on the left shows the "main path" through the network. The image on the right adds a shortcut to the main path. By stacking these ResNet blocks on top of each other, you can form a very deep network.
#
# The lecture mentioned that having ResNet blocks with the shortcut also makes it very easy for one of the blocks to learn an identity function. This means that you can stack on additional ResNet blocks with little risk of harming training set performance.
#
# On that note, there is also some evidence that the ease of learning an identity function accounts for ResNets' remarkable performance even more than skip connections help with vanishing gradients.
#
# Two main types of blocks are used in a ResNet, depending mainly on whether the input/output dimensions are the same or different. You are going to implement both of them: the "identity block" and the "convolutional block."

#
# ### 3.1 - The Identity Block
#
# The identity block is the standard block used in ResNets, and corresponds to the case where the input activation (say $a^{[l]}$) has the same dimension as the output activation (say $a^{[l+2]}$). To flesh out the different steps of what happens in a ResNet's identity block, here is an alternative diagram showing the individual steps:
#
#
#
   Figure 3   : Identity block. Skip connection "skips over" 2 layers.

#
# The upper path is the "shortcut path." The lower path is the "main path." In this diagram, notice the CONV2D and ReLU steps in each layer. To speed up training, a BatchNorm step has been added. Don't worry about this being complicated to implement--you'll see that BatchNorm is just one line of code in Keras!
#
# In this exercise, you'll actually implement a slightly more powerful version of this identity block, in which the skip connection "skips over" 3 hidden layers rather than 2 layers. It looks like this:
#
#
#
   Figure 4   : Identity block. Skip connection "skips over" 3 layers.

# These are the individual steps:
#
# First component of main path:
# - The first CONV2D has $F_1$ filters of shape (1,1) and a stride of (1,1). Its padding is "valid". Use 0 as the seed for the random uniform initialization: `kernel_initializer = initializer(seed=0)`.
# - The first BatchNorm is normalizing the 'channels' axis.
# - Then apply the ReLU activation function. This has no hyperparameters.
#
# Second component of main path:
# - The second CONV2D has $F_2$ filters of shape $(f,f)$ and a stride of (1,1). Its padding is "same". Use 0 as the seed for the random uniform initialization: `kernel_initializer = initializer(seed=0)`.
# - The second BatchNorm is normalizing the 'channels' axis.
# - Then apply the ReLU activation function. This has no hyperparameters.
#
# Third component of main path:
# - The third CONV2D has $F_3$ filters of shape (1,1) and a stride of (1,1). Its padding is "valid". Use 0 as the seed for the random uniform initialization: `kernel_initializer = initializer(seed=0)`.
# - The third BatchNorm is normalizing the 'channels' axis.
# - Note that there is **no** ReLU activation function in this component.
#
# Final step:
# - The `X_shortcut` and the output from the 3rd layer `X` are added together.
# - **Hint**: The syntax will look something like `Add()([var1,var2])`
# - Then apply the ReLU activation function. This has no hyperparameters.
#
#
# ### Exercise 1 - identity_block
#
# Implement the ResNet identity block. The first component of the main path has been implemented for you already! First, you should read these docs carefully to make sure you understand what's happening. Then, implement the rest.
# - To implement the Conv2D step: [Conv2D](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D)
# - To implement BatchNorm: [BatchNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization) `BatchNormalization(axis = 3)(X, training = training)`. If training is set to False, its weights are not updated with the new examples. I.e when the model is used in prediction mode.
# - For the activation, use:  `Activation('relu')(X)`
# - To add the value passed forward by the shortcut: [Add](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Add)
#
# We have added the initializer argument to our functions. This parameter receives an initializer function like the ones included in the package [tensorflow.keras.initializers](https://www.tensorflow.org/api_docs/python/tf/keras/initializers) or any other custom initializer. By default it will be set to [random_uniform](https://www.tensorflow.org/api_docs/python/tf/keras/initializers/RandomUniform)
#
# Remember that these functions accept a `seed` argument that can be any value you want, but that in this notebook must set to 0 for **grading purposes**.

#  Here is where you're actually using the power of the Functional API to create a shortcut path:

# In[2]:

# UNQ_C1
# GRADED FUNCTION: identity_block

def identity_block(X, f, filters, training=True, initializer=random_uniform):
    """
    Implementation of the identity block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    training -- True: Behave in training mode
                False: Behave in inference mode
    initializer -- to set up the initial weights of a layer. Equals to random uniform initializer

    Returns:
    X -- output of the identity block, tensor of shape (m, n_H, n_W, n_C)
    """

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value. You'll need this later to add back to the main path.
    X_shortcut = X

    # First component of main path
    X = Conv2D(filters = F1, kernel_size = 1, strides = (1,1), padding = 'valid', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training = training) # Default axis
    X = Activation('relu')(X)

    ### START CODE HERE
    ## Second component of main path (≈3 lines)
    X = Conv2D(filters = F2, kernel_size = (f, f), strides = (1,1), padding = 'same', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training = training)
    X = Activation('relu')(X)

    ## Third component of main path (≈2 lines)
    X = Conv2D(filters = F3, kernel_size = (1, 1), strides = (1,1), padding = 'valid', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training = training)

    ## Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)
    ### END CODE HERE

    return X

# In[3]:

np.random.seed(1)
X1 = np.ones((1, 4, 4, 3)) * -1
X2 = np.ones((1, 4, 4, 3)) * 1
X3 = np.ones((1, 4, 4, 3)) * 3

X = np.concatenate((X1, X2, X3), axis = 0).astype(np.float32)

A3 = identity_block(X, f=2, filters=[4, 4, 3],
                   initializer=lambda seed=0:constant(value=1),
                   training=False)
print('\033[1mWith training=False\033[0m\n')
A3np = A3.numpy()
print(np.around(A3.numpy()[:,(0,-1),:,:].mean(axis = 3), 5))
resume = A3np[:,(0,-1),:,:].mean(axis = 3)
print(resume[1, 1, 0])

print('\n\033[1mWith training=True\033[0m\n')
np.random.seed(1)
A4 = identity_block(X, f=2, filters=[3, 3, 3],
                   initializer=lambda seed=0:constant(value=1),
                   training=True)
print(np.around(A4.numpy()[:,(0,-1),:,:].mean(axis = 3), 5))

public_tests.identity_block_test(identity_block)

# **Expected value**
#
# ```
# With training=False
#
# [[[  0.        0.        0.        0.     ]
#   [  0.        0.        0.        0.     ]]
#
#  [[192.71234 192.71234 192.71234  96.85617]
#   [ 96.85617  96.85617  96.85617  48.92808]]
#
#  [[578.1371  578.1371  578.1371  290.5685 ]
#   [290.5685  290.5685  290.5685  146.78426]]]
# 96.85617
#
# With training=True
#
# [[[0.      0.      0.      0.     ]
#   [0.      0.      0.      0.     ]]
#
#  [[0.40739 0.40739 0.40739 0.40739]
#   [0.40739 0.40739 0.40739 0.40739]]
#
#  [[4.99991 4.99991 4.99991 3.25948]
#   [3.25948 3.25948 3.25948 2.40739]]]
# ```

#
# ### 3.2 - The Convolutional Block
#
# The ResNet "convolutional block" is the second block type. You can use this type of block when the input and output dimensions don't match up. The difference with the identity block is that there is a CONV2D layer in the shortcut path:
#
#
#
   Figure 4   : Convolutional block

#
# * The CONV2D layer in the shortcut path is used to resize the input $x$ to a different dimension, so that the dimensions match up in the final addition needed to add the shortcut value back to the main path. (This plays a similar role as the matrix $W_s$ discussed in lecture.)
# * For example, to reduce the activation dimensions's height and width by a factor of 2, you can use a 1x1 convolution with a stride of 2.
# * The CONV2D layer on the shortcut path does not use any non-linear activation function. Its main role is to just apply a (learned) linear function that reduces the dimension of the input, so that the dimensions match up for the later addition step.
# * As for the previous exercise, the additional `initializer` argument is required for grading purposes, and it has been set by default to [glorot_uniform](https://www.tensorflow.org/api_docs/python/tf/keras/initializers/GlorotUniform)
#
# The details of the convolutional block are as follows.
#
# First component of main path:
# - The first CONV2D has $F_1$ filters of shape (1,1) and a stride of (s,s). Its padding is "valid". Use 0 as the `glorot_uniform` seed `kernel_initializer = initializer(seed=0)`.
# - The first BatchNorm is normalizing the 'channels' axis.
# - Then apply the ReLU activation function. This has no hyperparameters.
#
# Second component of main path:
# - The second CONV2D has $F_2$ filters of shape (f,f) and a stride of (1,1). Its padding is "same".  Use 0 as the `glorot_uniform` seed `kernel_initializer = initializer(seed=0)`.
# - The second BatchNorm is normalizing the 'channels' axis.
# - Then apply the ReLU activation function. This has no hyperparameters.
#
# Third component of main path:
# - The third CONV2D has $F_3$ filters of shape (1,1) and a stride of (1,1). Its padding is "valid".  Use 0 as the `glorot_uniform` seed `kernel_initializer = initializer(seed=0)`.
# - The third BatchNorm is normalizing the 'channels' axis. Note that there is no ReLU activation function in this component.
#
# Shortcut path:
# - The CONV2D has $F_3$ filters of shape (1,1) and a stride of (s,s). Its padding is "valid".  Use 0 as the `glorot_uniform` seed `kernel_initializer = initializer(seed=0)`.
# - The BatchNorm is normalizing the 'channels' axis.
#
# Final step:
# - The shortcut and the main path values are added together.
# - Then apply the ReLU activation function. This has no hyperparameters.
#
#
# ### Exercise 2 - convolutional_block
#
# Implement the convolutional block. The first component of the main path is already implemented; then it's your turn to implement the rest! As before, always use 0 as the seed for the random initialization, to ensure consistency with the grader.
# - [Conv2D](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D)
# - [BatchNormalization](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization) (axis: Integer, the axis that should be normalized (typically the features axis)) `BatchNormalization(axis = 3)(X, training = training)`. If training is set to False, its weights are not updated with the new examples. I.e when the model is used in prediction mode.
# - For the activation, use:  `Activation('relu')(X)`
# - [Add](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Add)
#
# We have added the initializer argument to our functions. This parameter receives an initializer function like the ones included in the package [tensorflow.keras.initializers](https://www.tensorflow.org/api_docs/python/tf/keras/initializers) or any other custom initializer. By default it will be set to [random_uniform](https://www.tensorflow.org/api_docs/python/tf/keras/initializers/RandomUniform)
#
# Remember that these functions accept a `seed` argument that can be any value you want, but that in this notebook must set to 0 for **grading purposes**.

# In[4]:

# UNQ_C2
# GRADED FUNCTION: convolutional_block

def convolutional_block(X, f, filters, s = 2, training=True, initializer=glorot_uniform):
    """
    Implementation of the convolutional block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    s -- Integer, specifying the stride to be used
    training -- True: Behave in training mode
                False: Behave in inference mode
    initializer -- to set up the initial weights of a layer. Equals to Glorot uniform initializer,
                   also called Xavier uniform initializer.

    Returns:
    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
    """

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value
    X_shortcut = X

    ##### MAIN PATH #####

    # First component of main path glorot_uniform(seed=0)
    X = Conv2D(filters = F1, kernel_size = 1, strides = (s, s), padding='valid', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training=training)
    X = Activation('relu')(X)

    ### START CODE HERE

    ## Second component of main path (≈3 lines)
    X = Conv2D(filters = F2, kernel_size = (f,f), strides = (1, 1), padding='same', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training=training)
    X = Activation('relu')(X)

    ## Third component of main path (≈2 lines)
    X = Conv2D(filters = F3, kernel_size = (1,1), strides = (1, 1), padding='valid', kernel_initializer = initializer(seed=0))(X)
    X = BatchNormalization(axis = 3)(X, training=training)

    ##### SHORTCUT PATH ##### (≈2 lines)
    X_shortcut = Conv2D(filters = F3, kernel_size = (1,1), strides = (s, s), padding='valid', kernel_initializer = initializer(seed=0))(X_shortcut)
    X_shortcut = BatchNormalization(axis = 3)(X_shortcut, training=training)

    ### END CODE HERE

    # Final step: Add shortcut value to main path (Use this order [X, X_shortcut]), and pass it through a RELU activation
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X

# In[5]:

from outputs import convolutional_block_output1, convolutional_block_output2
np.random.seed(1)
#X = np.random.randn(3, 4, 4, 6).astype(np.float32)
X1 = np.ones((1, 4, 4, 3)) * -1
X2 = np.ones((1, 4, 4, 3)) * 1
X3 = np.ones((1, 4, 4, 3)) * 3

X = np.concatenate((X1, X2, X3), axis = 0).astype(np.float32)

A = convolutional_block(X, f = 2, filters = [2, 4, 6], training=False)

assert type(A) == EagerTensor, "Use only tensorflow and keras functions"
assert tuple(tf.shape(A).numpy()) == (3, 2, 2, 6), "Wrong shape."
assert np.allclose(A.numpy(), convolutional_block_output1), "Wrong values when training=False."
print(A[0])

B = convolutional_block(X, f = 2, filters = [2, 4, 6], training=True)
assert np.allclose(B.numpy(), convolutional_block_output2), "Wrong values when training=True."

print('\033[92mAll tests passed!')

# **Expected value**
#
# ```
# tf.Tensor(
# [[[0.         0.66683817 0.         0.         0.88853896 0.5274254 ]
#   [0.         0.65053666 0.         0.         0.89592844 0.49965227]]
#
#  [[0.         0.6312079  0.         0.         0.8636247  0.47643146]
#   [0.         0.5688321  0.         0.         0.85534114 0.41709304]]], shape=(2, 2, 6), dtype=float32)
# ```

#
# ## 4 - Building Your First ResNet Model (50 layers)
#
# You now have the necessary blocks to build a very deep ResNet. The following figure describes in detail the architecture of this neural network. "ID BLOCK" in the diagram stands for "Identity block," and "ID BLOCK x3" means you should stack 3 identity blocks together.
#
#
#
   Figure 5   : ResNet-50 model

#
# The details of this ResNet-50 model are:
# - Zero-padding pads the input with a pad of (3,3)
# - Stage 1:
#     - The 2D Convolution has 64 filters of shape (7,7) and uses a stride of (2,2).
#     - BatchNorm is applied to the 'channels' axis of the input.
#     - ReLU activation is applied.
#     - MaxPooling uses a (3,3) window and a (2,2) stride.
# - Stage 2:
#     - The convolutional block uses three sets of filters of size [64,64,256], "f" is 3, and "s" is 1.
#     - The 2 identity blocks use three sets of filters of size [64,64,256], and "f" is 3.
# - Stage 3:
#     - The convolutional block uses three sets of filters of size [128,128,512], "f" is 3 and "s" is 2.
#     - The 3 identity blocks use three sets of filters of size [128,128,512] and "f" is 3.
# - Stage 4:
#     - The convolutional block uses three sets of filters of size [256, 256, 1024], "f" is 3 and "s" is 2.
#     - The 5 identity blocks use three sets of filters of size [256, 256, 1024] and "f" is 3.
# - Stage 5:
#     - The convolutional block uses three sets of filters of size [512, 512, 2048], "f" is 3 and "s" is 2.
#     - The 2 identity blocks use three sets of filters of size [512, 512, 2048] and "f" is 3.
# - The 2D Average Pooling uses a window of shape (2,2).
# - The 'flatten' layer doesn't have any hyperparameters.
# - The Fully Connected (Dense) layer reduces its input to the number of classes using a softmax activation.
#
#
#
# ### Exercise 3 - ResNet50
#
# Implement the ResNet with 50 layers described in the figure above. We have implemented Stages 1 and 2. Please implement the rest. (The syntax for implementing Stages 3-5 should be quite similar to that of Stage 2) Make sure you follow the naming convention in the text above.
#
# You'll need to use this function:
# - Average pooling [see reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/AveragePooling2D)
#
# Here are some other functions we used in the code below:
# - Conv2D: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2D)
# - BatchNorm: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization) (axis: Integer, the axis that should be normalized (typically the features axis))
# - Zero padding: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/ZeroPadding2D)
# - Max pooling: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/MaxPool2D)
# - Fully connected layer: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense)
# - Addition: [See reference](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Add)

# In[9]:

# UNQ_C3
# GRADED FUNCTION: ResNet50

def ResNet50(input_shape = (64, 64, 3), classes = 6):
    """
    Stage-wise implementation of the architecture of the popular ResNet50:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> FLATTEN -> DENSE

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """

    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)

    # Zero-Padding
    X = ZeroPadding2D((3, 3))(X_input)

    # Stage 1
    X = Conv2D(64, (7, 7), strides = (2, 2), kernel_initializer = glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis = 3)(X)
    X = Activation('relu')(X)
    X = MaxPooling2D((3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, f = 3, filters = [64, 64, 256], s = 1)
    X = identity_block(X, 3, [64, 64, 256])
    X = identity_block(X, 3, [64, 64, 256])

    ### START CODE HERE

    ## Stage 3 (≈4 lines)
    X = convolutional_block(X, f = 3, filters = [128,128,512], s = 2)
    X = identity_block(X, 3, [128,128,512])
    X = identity_block(X, 3, [128,128,512])
    X = identity_block(X, 3, [128,128,512])

    ## Stage 4 (≈6 lines)
    X = convolutional_block(X, f = 3, filters = [256,256,1024], s = 2)
    X = identity_block(X, 3, [256, 256, 1024])
    X = identity_block(X, 3, [256, 256, 1024])
    X = identity_block(X, 3, [256, 256, 1024])
    X = identity_block(X, 3, [256, 256, 1024])
    X = identity_block(X, 3, [256, 256, 1024])

    ## Stage 5 (≈3 lines)
    X = convolutional_block(X, f = 3, filters = [512, 512, 2048], s = 2)
    X = identity_block(X, 3, [512, 512, 2048])
    X = identity_block(X, 3, [512, 512, 2048])

    ## AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    X = X = AveragePooling2D((2,2))(X)

    ### END CODE HERE

    # output layer
    X = Flatten()(X)
    X = Dense(classes, activation='softmax', kernel_initializer = glorot_uniform(seed=0))(X)

    # Create model
    model = Model(inputs = X_input, outputs = X)

    return model

# Run the following code to build the model's graph. If your implementation is incorrect, you'll know it by checking your accuracy when running `model.fit(...)` below.

# In[10]:

model = ResNet50(input_shape = (64, 64, 3), classes = 6)
print(model.summary())

# In[11]:

from outputs import ResNet50_summary

model = ResNet50(input_shape = (64, 64, 3), classes = 6)

comparator(summary(model), ResNet50_summary)

# As shown in the Keras Tutorial Notebook, prior to training a model, you need to configure the learning process by compiling the model.

# In[12]:

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# The model is now ready to be trained. The only thing you need now is a dataset!

# Let's load your old friend, the SIGNS dataset.
#
#
#
   Figure 6   : SIGNS dataset

#

# In[13]:

X_train_orig, Y_train_orig, X_test_orig, Y_test_orig, classes = load_dataset()

# Normalize image vectors
X_train = X_train_orig / 255.
X_test = X_test_orig / 255.

# Convert training and test labels to one hot matrices
Y_train = convert_to_one_hot(Y_train_orig, 6).T
Y_test = convert_to_one_hot(Y_test_orig, 6).T

print ("number of training examples = " + str(X_train.shape[0]))
print ("number of test examples = " + str(X_test.shape[0]))
print ("X_train shape: " + str(X_train.shape))
print ("Y_train shape: " + str(Y_train.shape))
print ("X_test shape: " + str(X_test.shape))
print ("Y_test shape: " + str(Y_test.shape))

# Run the following cell to train your model on 10 epochs with a batch size of 32. On a GPU, it should take less than 2 minutes.

# In[14]:

model.fit(X_train, Y_train, epochs = 10, batch_size = 32)

# **Expected Output**:
#
# ```
# Epoch 1/10
# 34/34 [==============================] - 1s 34ms/step - loss: 1.9241 - accuracy: 0.4620
# Epoch 2/10
# 34/34 [==============================] - 2s 57ms/step - loss: 0.6403 - accuracy: 0.7898
# Epoch 3/10
# 34/34 [==============================] - 1s 24ms/step - loss: 0.3744 - accuracy: 0.8731
# Epoch 4/10
# 34/34 [==============================] - 2s 44ms/step - loss: 0.2220 - accuracy: 0.9231
# Epoch 5/10
# 34/34 [==============================] - 2s 57ms/step - loss: 0.1333 - accuracy: 0.9583
# Epoch 6/10
# 34/34 [==============================] - 2s 52ms/step - loss: 0.2243 - accuracy: 0.9444
# Epoch 7/10
# 34/34 [==============================] - 2s 48ms/step - loss: 0.2913 - accuracy: 0.9102
# Epoch 8/10
# 34/34 [==============================] - 1s 30ms/step - loss: 0.2269 - accuracy: 0.9306
# Epoch 9/10
# 34/34 [==============================] - 2s 46ms/step - loss: 0.1113 - accuracy: 0.9630
# Epoch 10/10
# 34/34 [==============================] - 2s 57ms/step - loss: 0.0709 - accuracy: 0.9778
# ```
#
# The exact values could not match, but don't worry about that. The important thing that you must see is that the loss value decreases, and the accuracy increases for the firsts 5 epochs.

# Let's see how this model (trained on only two epochs) performs on the test set.

# In[15]:

preds = model.evaluate(X_test, Y_test)
print ("Loss = " + str(preds[0]))
print ("Test Accuracy = " + str(preds[1]))

# **Expected Output**:
#
#

#

#
#             Test Accuracy
#
#
#            >0.80
#
#

#
#

# For the purposes of this assignment, you've been asked to train the model for ten epochs. You can see that it performs well. The online grader will only run your code for a small number of epochs as well. Please go ahead and submit your assignment.

# After you have finished this official (graded) part of this assignment, you can also optionally train the ResNet for more iterations, if you want. It tends to get much better performance when trained for ~20 epochs, but this does take more than an hour when training on a CPU.
#
# Using a GPU, this ResNet50 model's weights were trained on the SIGNS dataset. You can load and run the trained model on the test set in the cells below. It may take ≈1min to load the model. Have fun!

# In[16]:

pre_trained_model = tf.keras.models.load_model('resnet50.h5')

# In[17]:

preds = pre_trained_model.evaluate(X_test, Y_test)
print ("Loss = " + str(preds[0]))
print ("Test Accuracy = " + str(preds[1]))

# **Congratulations** on finishing this assignment! You've now implemented a state-of-the-art image classification system! Woo hoo!
#
# ResNet50 is a powerful model for image classification when it's trained for an adequate number of iterations. Hopefully, from this point, you can use what you've learned and apply it to your own classification problem to perform state-of-the-art accuracy.

#
#
# **What you should remember**:
#
# - Very deep "plain" networks don't work in practice because vanishing gradients make them hard to train.
# - Skip connections help address the Vanishing Gradient problem. They also make it easy for a ResNet block to learn an identity function.
# - There are two main types of blocks: The **identity block** and the **convolutional block**.
# - Very deep Residual Networks are built by stacking these blocks together.

#
# ## 5 - Test on Your Own Image (Optional/Ungraded)

# If you wish, you can also take a picture of your own hand and see the output of the model. To do this:
#     1. Click on "File" in the upper bar of this notebook, then click "Open" to go on your Coursera Hub.
#     2. Add your image to this Jupyter Notebook's directory, in the "images" folder
#     3. Write your image's name in the following code
#     4. Run the code and check if the algorithm is right!

# In[ ]:

img_path = 'images/my_image.jpg'
img = image.load_img(img_path, target_size=(64, 64))
x = image.img_to_array(img)
x = np.expand_dims(x, axis=0)
x = x/255.0
print('Input image shape:', x.shape)
imshow(img)
prediction = pre_trained_model.predict(x)
print("Class prediction vector [p(0), p(1), p(2), p(3), p(4), p(5)] = ", prediction)
print("Class:", np.argmax(prediction))

# Even though the model has high accuracy, it might be performing poorly on your own set of images. Notice that, the shape of the pictures, the lighting where the photos were taken, and all of the preprocessing steps can have an impact on the performance of the model. Considering everything you have learned in this specialization so far, what do you think might be the cause here?
#
# *Hint*: It might be related to some distributions. Can you come up with a potential solution ?

# You can also print a summary of your model by running the following code.

# In[ ]:

pre_trained_model.summary()

#
# ## 6 - Bibliography
#
# This notebook presents the ResNet algorithm from He et al. (2015). The implementation here also took significant inspiration and follows the structure given in the GitHub repository of Francois Chollet:
#
# - Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun - [Deep Residual Learning for Image Recognition (2015)](https://arxiv.org/abs/1512.03385)
# - Francois Chollet's GitHub repository: https://github.com/fchollet/deep-learning-models/blob/master/resnet50.py
#

# In[ ]:

4 - Convolutional Neural Networks/Residual_Networks.py


import numpy as np
import torch
import torch.utils.data as data
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import os, random, glob, cv2

class DALY2G(data.Dataset):
    def __init__(self, data_dir='data', text_file='test_img.txt', transform=None):
        super(DALY2G, self).__init__()
        # load split file
        text_path = os.path.join(data_dir, text_file)
        self.indexlist = [line.rstrip('\n') for line in open(text_path,'r')]

        # load images
        self.data_dir = data_dir
        self.img_dir = 'images_w_act_raw'
        self.transform = transform

    def load_img(self, index):
        info = self.indexlist[index]
        vid, im_name = info.split('/')
        timestamp = im_name.split('_')[1]

        img_path = os.path.join(self.data_dir, self.img_dir, info)
        img = Image.open(img_path).convert('RGB')
        w, h = img.size

        annot_path = os.path.join(self.data_dir, 'mtcnn_bbox_kpnts', vid, 'img_'+timestamp+'.npz')
        annot = np.load(annot_path)

        face_bbox = torch.DoubleTensor(40,5).zero_()
        if len(annot['gt']) > 0:
            #  print('len:', len(annot['gt']))
            for i in range(len(annot['gt'])):
                #  print('i:' , i)
                _bbox = annot['gt'][i]
                # enlarge the bbox 2.2x
                x1 = _bbox[0]
                y1 = _bbox[1]
                x2 = _bbox[2]
                y2 = _bbox[3]
                _bbox[0] = max(0, 1.6*x1 - 0.6*x2)
                _bbox[1] = max(0, 1.6*y1 - 0.6*y2)
                _bbox[2] = min(w-1, 1.6*x2 - 0.6*x1)
                _bbox[3] = min(h-1, 1.6*y2 - 0.6*y1)
                face_bbox[i] = torch.DoubleTensor(_bbox)
            face_count = len(annot['gt'])
        else:
            face_count = 0

        return img, face_bbox, face_count

    def __getitem__(self, index):
        _img, face_bbox, face_count = self.load_img(index)

        # transform images if required
        if self.transform:
            img = self.transform(_img)

        #  print(img.size(), face_bbox.size(), face_count)
        return img, face_bbox, face_count, index

    def __len__(self):
        return len(self.indexlist)

def main():
    import torch
    from torchvision import transforms, utils
    from torch.utils.data import DataLoader
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    g_data = DALY2G(data_dir='/home/SSD5/jason-data/Privacy/dataset/DALY',
            text_file='test_img.txt',
            transform=transforms.Compose([
                        transforms.Resize((600, 800)),
                        transforms.ToTensor(), normalize]))
    print(len(g_data))
    dataloader = DataLoader(g_data, batch_size=4, shuffle=False, num_workers=1)
    for i, data in enumerate(dataloader):
        im, bbox, flag = data
        print(im.size(), bbox.size(), flag.size())
        utils.save_image(im, str(i) + '.jpg', normalize=True)

if __name__ == "__main__":
    main()


#!/usr/bin/env python
#
# // SPDX-License-Identifier: BSD-3-CLAUSE
#
# (C) Copyright 2018, Xilinx, Inc.
#
import numpy as np
# from . import xdnn
import xdnn
import argparse
import os
import json
import math
import cv2
import PyTurboJPEG

def default_parser_args():
    parser = argparse.ArgumentParser(description='pyXDNN')
    parser.add_argument('--xdnnv3', type = str, default = 'false',
        help='is xdnnv3')
    parser.add_argument('--platform',
        help='What platform are you running on')
    parser.add_argument('--xclbin',
        help='.xclbin file')
    parser.add_argument('--batch_sz', type=int, default=1,
        help='batch size')
    parser.add_argument('--netcfg',
        help='FPGA instructions generated by compiler for the network')
    parser.add_argument('--quantizecfg', default="",
        help='FPGA config file')
    parser.add_argument('--xlnxlib',
        help='FPGA xfDNN lib .so')
    parser.add_argument('--fpgaoutsz', type=int, default=1024,
        help='size of 1 FPGA output blob')
    parser.add_argument('--outsz', type=int, default=1000,
        help='size of last layer\'s output blob')
    parser.add_argument('--firstfpgalayer', default="",
        help='name of first FPGA layer (to start quantization)')
    parser.add_argument('--datadir',
        help='path to data files to run for the network')
    parser.add_argument('--labels', default="",
        help='result -> labels translation file')
    parser.add_argument('--jsoncfg',
        help='json file with nets, data and PEs to use'),
    parser.add_argument('--images', nargs='*',
        help='raw image files to use as input')
    parser.add_argument('--transform', type=str, default='resize',
        help='input image processing transform (i.e. resize, yolo, etc.')
    parser.add_argument('--scaleA', type=int, default=10000,
        help='weights scaling value')
    parser.add_argument('--scaleB', type=int, default=30,
        help='activation scaling value ')
    parser.add_argument('--img_raw_scale', type=float, default=255.0,
        help='image raw scale value ')
    parser.add_argument('--img_mean', type=tuple, default=(104.007, 116.669, 122.679),  # BGR for Caffe
        help='image mean values ')
    parser.add_argument('--img_input_scale', type=float, default=1.0,
        help='image input scale value ')
    parser.add_argument('--in_shape', type=tuple, default=(3, 224, 224), help='input dimensions')
    parser.add_argument('--useblas', default=False, action='store_true',
        help='use BLAS-optimized functions (requires xfDNN lib compiled with BLAS)')
    parser.add_argument('--PE', nargs='?', type=int, default=-1,
        help='preferred PE to run the classification on. Default is auto-select')
    parser.add_argument('--blocking', nargs='?',
            help='use blocking calls. Default is blocking')
    parser.add_argument('--endLayerName', default="",
            help='layer name till the network should be run, helpful for debugging')
    parser.add_argument('--diffStartLayer', type=int, default=0,
            help="if 1 then we can run from any given layer ignoring the X's of first layers")
    parser.add_argument('--v2WeightsFormat', type=int, default=0,
            help="Weights File specified as KernSizex KernSizey instead of only KernSize, supporting rectangle kernels")
    parser.add_argument('--layerName', default="",
            help='layername until which pyfpga should run, if left default, would run the entire model')
    parser.add_argument('--binaryFormatWeights', type=int, default=0,
            help="Binary Format Weights Files")
    return parser

def make_dict_args(args):
    args_dict = vars(args)

    if args_dict['images'] is not None and os.path.isdir(args_dict['images'][0]):
        inputFiles = [os.path.join(args_dict['images'][0], f) \
        for f in os.listdir(args_dict['images'][0]) if os.path.isfile(os.path.join(args_dict['images'][0], f))]
        args_dict['images'] = inputFiles

    if args_dict['jsoncfg']:
            fileName = args_dict['jsoncfg']
            args_dict['jsoncfg'] = args_dict
            with open(fileName) as jsoncfgFile:
                    args_dict['jsoncfg'] = json.load(jsoncfgFile)['confs']
                    new_args = []
                    for dict_elem in args_dict['jsoncfg']:
                            # if item.key
                            # for key,value in dict_elem.items():
                            for key, value in list(args_dict.items()):
                                    if key not in dict_elem:
                                            dict_elem[key] = value

                            new_args.append(dict_elem)

                    # print new_args
                    args_dict['jsoncfg'] = new_args
                    for dict_elem in args_dict['jsoncfg']:
                      for key,val in dict_elem.iteritems():
                        try:
                          dict_elem[key] = int(val)
                        except:
                          pass
    return args_dict

def processCommandLine(argv=None):
    """
    Invoke command line parser for command line deployment flows.
    """
    parser = default_parser_args()
    args = parser.parse_args(argv)
    return make_dict_args(args)

def loadImageBlobFromFile(imgFile, raw_scale, mean, input_scale, img_h, img_w):
    if isinstance(imgFile,str):
        img = PyTurboJPEG.imread  ( imgFile)

    height, width, _ = img.shape
    if height != img_h or width != img_w:
        img = cv2.resize(img, (img_h, img_w))

    #print "Bryan: ", np.max(img)
    #print "Bryan: ", np.min(img)

    # Not needed for Caffe because BGR:
    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # FIXME TF retrained
    if raw_scale != 255:
        array = img.astype(np.float32) / 255
        array *= raw_scale
    mean_arr = np.array(mean)
    array = img - mean_arr
    if input_scale != 1:
        array *= input_scale

    # prepare blob
    array = np.transpose(array, (2, 0, 1))
    array = np.broadcast_to( array, (1,) + array.shape )

    # array = array.flatten() # FIXME TF retrained

    return array

def loadYoloImageBlobFromFile(imgFile, img_h, img_w):
    # This first loads the image, runs BGR->RGB, converts int 0-255 to floats 0.0-1.0, then letterboxes

    import cv2
    img = cv2.imread(imgFile)

    # YOLO was trained with RGB, not BGR like Caffe:
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    height, width, channels = img.shape

    newdim = max(height, width)
    scalew = float(width) / newdim
    scaleh = float(height) / newdim
    neww = int(img_w * scalew)
    newh = int(img_h * scaleh)

    reduced_img = cv2.resize(img, (neww, newh))

    reduced_img_float = reduced_img / 255.0

    newdim = max(newh, neww)
    diffh = newdim - newh
    diffw = newdim - neww
    letter_image_float = np.zeros((newdim, newdim, channels), np.float)
    letter_image_float[:, :, :] = 0.5
    letter_image_float[diffh / 2:newh + diffh / 2, diffw / 2:neww + diffw / 2] = reduced_img_float

    # prepare blob (channel, row, column)
    array = np.transpose(letter_image_float, (2, 0, 1))
    array = np.broadcast_to( array, (1,) + array.shape )
    #array = np.transpose(letter_image_float, (2, 0, 1)).flatten()

    return array, img.shape

def loadImages (files, transform, img_raw_scale, img_mean, img_input_scale, img_shape):
    inputs = np.zeros( ( (len(files),) + img_shape), dtype=np.float32)

    shapes = [None] * len(files)
    # use raw image files from user

    for i, fname in enumerate(files):
      if transform == 'resize':
        inputs[i] = loadImageBlobFromFile(fname, img_raw_scale, img_mean, img_input_scale, img_shape[1], img_shape[2])
        shapes[i] = None
      elif transform == 'yolo':
        inputs[i], shapes[i] = loadYoloImageBlobFromFile(fname, img_shape[1], img_shape[2])
      else:
        assert 0, "Unrecognized image transform"

    return inputs, shapes

def loadRawDataInput (file, num_images, img_shape):
    import os
    from collections import Iterable
    if not isinstance(img_shape, Iterable):
        img_shape = eval(img_shape)

    inputs = np.zeros((num_images, np.prod(img_shape)), dtype=np.float32)
    shapes = [None] * num_images

# use saved image data from framework
    fname = "%s/input" % file
    if not os.path.isfile(fname):
        raise Exception('Input file not found')

    with open(fname, 'r') as f:
        for i in range(num_images):
                line = f.readline()
                if not line:
                    break

                vals = line.strip().split(' ')
                vals = [float(v) for v in vals]

                inputs[i] = vals
                shapes[i] = img_shape
        return inputs, shapes

def getClassification(output, args):
    if isinstance (output, np.ndarray):
        output = output.flatten().tolist()

    labels = []
    with open(args['labels'], 'r') as f:
        for line in f:
            labels.append(line.strip())

    idxArr   = [i for i in range(args['outsz'])]
    batch_sz = len(output) / args['outsz']
    result   = []
    for i in range(batch_sz):
      startIdx = i * args['outsz']
      vals     = output[startIdx:startIdx + args['outsz']]
      top5     = sorted(zip(vals, idxArr), reverse=True)[:5]
      result  += [[(top5[j][0], labels[top5[j][1]]) for j in range(len(top5))]]

    return result

def printClassification(output, args):
    """
    Print the result of classification given class scores, and a synset labels file.

    :param output: Class scores, typically the output of the softmax layer.
    :type output: numpy.ndarray.
    :param args: Collection of arguments. Most importanly args["labels"] which is the path to a file containing class names.
    :type args: dict.
    """
    top5s = getClassification(output, args)

    print("\n")
    for i, top5 in enumerate(top5s):
      inputImage = ""
      if args['images']:
        inputImage = "for %s " % args['images'][i]

      print("---------- Prediction %d/%d %s----------" % (i, len(top5s)-1, inputImage))

      for (prob, label) in top5:
        print("%.4f - \"%s\"" % (prob, label))
      print("")

def XDLFprepareRawInputs(args, RawInputs, PE=-1):
    fpgaInputs = xdnn.prepareInputsForFpga(RawInputs, args['quantizecfg'], args['scaleB'], PE, args['firstfpgalayer'])
    return fpgaInputs

def prepareOutput(output_sz, batch_sz):
  (fpgaOutput, _) = xdnn.makeFPGAFloatArray(output_sz*batch_sz)

  return fpgaOutput

def XDLFloadWeights ( args,Weights,outChans,inChans,kernH,kernW,layerName, isxdnnv3=False) :
  print "Loading weights/bias/quant_params to FPGA..."

  if isxdnnv3=="True":
    size = xdnn.v3computeWeightsBiasQuantSize(kernW, kernH, outChans, int(math.ceil(float(inChans)/float(96))), 0, 0, False)
    size=size*2
  else:
    size = xdnn.computeWeightsBiasQuantSize(\
          kernW,kernH,inChans,outChans,True if args['quantizecfg'] else False)

  blob = xdnn.makeWeightsBiasQuantBlob(size)
  bias=[0 for v in range(outChans)]
  if isxdnnv3=="True":
    offset = xdnn.v3fillWeightsBiasQuantBlob(blob, 0,
          args['quantizecfg'], Weights, args['scaleA'], bias, args['scaleB'],
          kernW,kernH,inChans,outChans,layerName)
  else:
    offset = xdnn.fillWeightsBiasQuantBlob(blob, 0,
          args['quantizecfg'], Weights, args['scaleA'], bias, args['scaleB'],
          kernW,kernH,inChans,outChans,layerName)
  fps = xdnn.loadBlobToDdr(blob, size, int(args['PE']))

  return (blob, fps)

def XDLFBunchComputeSizeLatestRepl(args, outChans, inChans, kernH, kernW,srcFullSectNum, srcReplSectNum, srcReplUnitNum, isxdnnv3=False):
  is8Bit=True
  if isxdnnv3=="True":
    size = xdnn.v3computeWeightsBiasQuantSize(kernW, kernH, outChans, int(srcFullSectNum), int(srcReplSectNum), int(srcReplUnitNum), True)
    size=size*2
  else:
    size = xdnn.computeWeightsBiasQuantSize(\
          kernW,kernH,inChans,outChans,True if args['quantizecfg'] else False)

  return size

def prepareRawInputs(args, PE=-1):
    print('import Raw Input')
    if args['batch_sz'] != 1:
        batch_sz = int(args['batch_sz'])
    else:
        batch_sz = 1
    inputs, shapes = loadRawDataInput (args['datadir'], batch_sz, args['in_shape'])
    #print('image shape {}, path {}, {}'.format(inputs.shape, args['datadir'], inputs[:10]))
    fpgaInputs = xdnn.prepareInputsForFpga(inputs, args['quantizecfg'], args['scaleB'], PE, args['firstfpgalayer'])

    return fpgaInputs, shapes, batch_sz

def prepareImages(args, PE=-1):
    batch_sz = len(args['images'])
    inputs, shapes = loadImages(args['images'], args['transform'], args['img_raw_scale'], args['img_mean'], args['img_input_scale'], args['in_shape'])
    #print('image shape {}, path {}, {}'.format(inputs.shape, args['images'], inputs[:10]))
    fpgaInputs = xdnn.prepareInputsForFpga(inputs, args['quantizecfg'], args['scaleB'], PE, args['firstfpgalayer'])

    return fpgaInputs, shapes, batch_sz

def loadWeights (args) :
    """
    Load weights to off chip device memory. The weights are first quantized.

    :param args: Collection of arguments. Most importanly args["datadir"] which is the path to a folder containing weights & biases.
    :type args: dict.
    :returns: tuple -- (weightsBlob, fcWeight, fcBias) -- <class 'xdnn.LP_c_short'>, numpy.ndarray, numpy.ndarray
    """
    if 'xdnnv3' in args and args['xdnnv3'] == 'true':
      weightsBlob = loadWeightsBiasQuantv3(args)
    else:
      weightsBlob = loadWeightsBiasQuant(args)
    (fcWeight, fcBias) = loadFCWeightsBias(args['datadir'])
    return (weightsBlob, fcWeight, fcBias)

def prepareInput (args, PE=-1):
    """
    Allocate memory for inputs, load images from disk, images will be scaled, letteboxed, and preprocessed.

    :param args: Collection of arguments. Most importanly args["images"] which is a list of images to prepare.
    :type args: dict.
    :param PE: Index used to target specific processing element. Use -1 for autoselect. There can be from 1 to 6 processing elements in a particular xclbin.
    :type PE: int.
    :returns: tuple -- (fpgaInputs,batch_sz) -- <class 'xdnn.LP_c_short_Array_1'>, int
    """
    if args['images'] is not None:
        (fpgaInputs, shapes, batch_sz) = prepareImages(args, PE)
    else:
        (fpgaInputs, shapes, batch_sz) = prepareRawInputs(args, PE)

    if args['transform'] == 'yolo':
        return fpgaInputs, shapes, batch_sz
    else:
        return fpgaInputs, batch_sz

def prepareOutput(output_sz, batch_sz):
    """
    Allocate memory for outputs.

    :param output_sz: Number of elements in the output volume returned by FPGA for a single inference. This is specific to the network.
    :type output_sz: int.
    :param batch_sz: Number of images to be processed simultaneously.
    :type batch_sz: int.
    :returns: numpy.ndarray -- 1D Array large enough to hold output_sz*batch_sz elements.
    """
    (fpgaOutput, fpgaHandle) = xdnn.makeFPGAFloatArray(output_sz * batch_sz)

    return fpgaOutput

def XDLFloadWeights (args, Weights, outChans, inChans, kernH, kernW, layerName, isxdnnv3=False) :
    print("Loading weights/bias/quant_params to FPGA...")

    if isxdnnv3 == "True":
        size = xdnn.v3computeWeightsBiasQuantSize(kernW, kernH, outChans, int(math.ceil(float(inChans) / float(96))), 0, 0, False)
        size = size * 2
    else:
        size = xdnn.computeWeightsBiasQuantSize(\
                    kernW, kernH, inChans, outChans, True if args['quantizecfg'] else False)

    blob = xdnn.makeWeightsBiasQuantBlob(size)
    bias = [0 for v in range(outChans)]
    if isxdnnv3 == "True":
        offset = xdnn.v3fillWeightsBiasQuantBlob(blob, 0,
                    args['quantizecfg'], Weights, args['scaleA'], bias, args['scaleB'],
                    kernW, kernH, inChans, outChans, layerName)
    else:
        offset = xdnn.fillWeightsBiasQuantBlob(blob, 0,
                    args['quantizecfg'], Weights, args['scaleA'], bias, args['scaleB'],
                    kernW, kernH, inChans, outChans, layerName)
    layer2OffsetMap = "%s:%d" % (layerName, 0)
    fps = xdnn.loadBlobToDdr(blob, size, layer2OffsetMap, int(args['PE']))

    return (blob, fps)

def XDLFBunchComputeSizeLatestRepl(args, outChans, inChans, kernH, kernW, srcFullSectNum, srcReplSectNum, srcReplUnitNum, isxdnnv3=False):

    if isxdnnv3 == "True":
        size = xdnn.v3computeWeightsBiasQuantSize(kernW, kernH, outChans, int(srcFullSectNum), int(srcReplSectNum), int(srcReplUnitNum), False)
        size = size * 2
    else:
        size = xdnn.computeWeightsBiasQuantSize(\
                    kernW, kernH, inChans, outChans, True if args['quantizecfg'] else False)

    return size

def XDLFBunchloadWeightsBiasQuantLatestRepl(args, allLayerNames, allLayersWeightsBiasQuantizeKey, allConvLayerNamesParams, size, isxdnnv3):
    print("Loading weights/bias/quant_params to FPGA...")

    blob = xdnn.makeWeightsBiasQuantBlob(size)

    offset = 0

    layer2OffsetMapStr = ""
    for layerName in allLayerNames:
        if layer2OffsetMapStr != "":
            layer2OffsetMapStr += ","
        layer2OffsetMapStr += "%s:%d" % (layerName, offset)

        if isxdnnv3 == "True":
            offset += xdnn.v3fillWeightsBiasQuantBlob(blob, offset,
                    args['quantizecfg'], allLayersWeightsBiasQuantizeKey[layerName]['weights'], args['scaleA'], allLayersWeightsBiasQuantizeKey[layerName]['Bias'], args['scaleB'],
                 allLayersWeightsBiasQuantizeKey[layerName]['kernW'], allLayersWeightsBiasQuantizeKey[layerName]['kernH'], allLayersWeightsBiasQuantizeKey[layerName]['inChans'], allLayersWeightsBiasQuantizeKey[layerName]['outChans'], int(allConvLayerNamesParams[layerName]['srcFullSectNum']), int(allConvLayerNamesParams[layerName]['srcReplSectNum']), int(allConvLayerNamesParams[layerName]['srcReplUnitNum']), int(allConvLayerNamesParams[layerName]['srcReplUnitWidth']), int(allConvLayerNamesParams[layerName]['convHalfRateMode']), layerName)
        else:
            offset += xdnn.fillWeightsBiasQuantBlob(blob, offset,
                    args['quantizecfg'], allLayersWeightsBiasQuantizeKey[layerName]['weights'], args['scaleA'], allLayersWeightsBiasQuantizeKey[layerName]['Bias'], args['scaleB'],
                 allLayersWeightsBiasQuantizeKey[layerName]['kernW'], allLayersWeightsBiasQuantizeKey[layerName]['kernH'], allLayersWeightsBiasQuantizeKey[layerName]['inChans'], allLayersWeightsBiasQuantizeKey[layerName]['outChans'], layerName)
    fps = xdnn.loadBlobToDdr(blob, size, layer2OffsetMapStr, int(args['PE']))

    return (blob, fps)

def XDLFBunchComputeSize(args, outChans, inChans, kernH, kernW, isxdnnv3=False):

    if isxdnnv3 == "True":
        size = xdnn.v3computeWeightsBiasQuantSize(kernW, kernH, outChans, int(math.ceil(float(inChans) / float(96))), 0, 0, False)
        size = size * 2
    else:
        size = xdnn.computeWeightsBiasQuantSize(\
                    kernW, kernH, inChans, outChans, True if args['quantizecfg'] else False)

    return size

def XDLFBunchloadWeightsBiasQuant(args, allLayerNames, allLayersWeightsBiasQuantizeKey, size, isxdnnv3):
    print("Loading weights/bias/quant_params to FPGA...")

    print("MNDBG total A size %d" % size)
    blob = xdnn.makeWeightsBiasQuantBlob(size)

    layer2OffsetMapStr = ""
    offset = 0

    for layerName in allLayerNames:
        if layer2OffsetMapStr != "":
            layer2OffsetMapStr += ","
        layer2OffsetMapStr += "%s:%d" % (layerName, offset)

        if isxdnnv3 == "True":
            offset += xdnn.v3fillWeightsBiasQuantBlob(blob, offset,
                    args['quantizecfg'], allLayersWeightsBiasQuantizeKey[layerName]['weights'], args['scaleA'], allLayersWeightsBiasQuantizeKey[layerName]['Bias'], args['scaleB'],
                 allLayersWeightsBiasQuantizeKey[layerName]['kernW'], allLayersWeightsBiasQuantizeKey[layerName]['kernH'], allLayersWeightsBiasQuantizeKey[layerName]['inChans'], allLayersWeightsBiasQuantizeKey[layerName]['outChans'], layerName)
        else:
            offset += xdnn.fillWeightsBiasQuantBlob(blob, offset,
                    args['quantizecfg'], allLayersWeightsBiasQuantizeKey[layerName]['weights'], args['scaleA'], allLayersWeightsBiasQuantizeKey[layerName]['Bias'], args['scaleB'],
                 allLayersWeightsBiasQuantizeKey[layerName]['kernW'], allLayersWeightsBiasQuantizeKey[layerName]['kernH'], allLayersWeightsBiasQuantizeKey[layerName]['inChans'], allLayersWeightsBiasQuantizeKey[layerName]['outChans'], layerName)
    fps = xdnn.loadBlobToDdr(blob, size, layer2OffsetMapStr, int(args['PE']))

    return (blob, fps)

def parseCompilerFile(compilerFileName):
  with open(compilerFileName) as compilerReadStream:
    compilerContent = compilerReadStream.readlines()
  compilerContent = [x.strip().split(" ") for x in compilerContent]
  allLayersParams={}
  layerParams={}
  allLayerNames=[]
  for i in range(len(compilerContent)):
    if compilerContent[i][1] == "XNConv":
      layerParams={}
      if compilerContent[i][2] not in allLayerNames:
        allLayerNames.append(compilerContent[i][2])
      layerParams['kernW']=compilerContent[i][3]
      layerParams['kernH']=compilerContent[i][4]
      layerParams['inChans']=compilerContent[i][19]
      layerParams['outChans']=compilerContent[i][23]
      layerParams['srcFullSectNum']=compilerContent[i][25]
      layerParams['srcReplSectNum']=compilerContent[i][26]
      layerParams['srcReplUnitNum']=compilerContent[i][27]
      layerParams['srcReplUnitWidth']=compilerContent[i][28]
      layerParams['convHalfRateMode']=compilerContent[i][47]

      allLayersParams[compilerContent[i][2]]=layerParams
    elif compilerContent[i][1] == "XNMaxPoolPipelined":
      layerParams={}
      if compilerContent[i][47] not in allLayerNames:
        allLayerNames.append(compilerContent[i][47])
      layerParams['kernW']=compilerContent[i][48]
      layerParams['kernH']=compilerContent[i][49]
      layerParams['inChans']=compilerContent[i][12]
      layerParams['outChans']=compilerContent[i][60]
      layerParams['srcFullSectNum']=compilerContent[i][17]
      layerParams['srcReplSectNum']=compilerContent[i][18]
      layerParams['srcReplUnitNum']=compilerContent[i][19]
      layerParams['srcReplUnitWidth']=compilerContent[i][20]
      layerParams['convHalfRateMode']=compilerContent[i][39]

      allLayersParams[compilerContent[i][47]]=layerParams

  return allLayerNames, allLayersParams

def loadWeightsBiasQuantv3(args):
    print("Loading weights/bias/quant_params to FPGA...")
    allConvLayerNames, allConvLayerNamesParams = parseCompilerFile(args['netcfg'])

    size = 0
    fi = 0
    while True:
        fname = "%s/fwbqb_%d" % (args['datadir'], fi)
        if not os.path.isfile(fname):
            break

        layerName = None
        weights = None
        bias = None
        kern = None
        inChans = None
        outChans = None

        with open(fname, 'r') as f:
            data = f.read()
            vals = data.strip().split(' ')
            layerName = vals[0]
            kern = int(vals[1])
            assert kern == int(allConvLayerNamesParams[layerName]['kernW'])
            assert kern == int(allConvLayerNamesParams[layerName]['kernH'])
            inChans = int(vals[2])
            assert inChans == int(allConvLayerNamesParams[layerName]['inChans'])
            outChans = int(vals[3])
    #        assert outChans == int(allConvLayerNamesParams[layerName]['outChans'])

            srcFullSectNum = int(allConvLayerNamesParams[layerName]['srcFullSectNum'])
            srcReplSectNum = int(allConvLayerNamesParams[layerName]['srcReplSectNum'])
            srcReplUnitNum = int(allConvLayerNamesParams[layerName]['srcReplUnitNum'])
        size += xdnn.v3computeWeightsBiasQuantSize(kern, kern, outChans, srcFullSectNum, srcReplSectNum, srcReplUnitNum, False)

        fi += 1

    size = size*2
    blob = xdnn.makeWeightsBiasQuantBlob(size)

    fi = 0
    offset = 0
    layer2OffsetMapStr = ""
    while True:
        fname = "%s/fwbqb_%d" % (args['datadir'], fi)
        if not os.path.isfile(fname):
            break

        layerName = None
        weights = None
        bias = None
        kern = None
        inChans = None
        outChans = None

        with open(fname, 'r') as f:
            data = f.read()
            vals = data.strip().split(' ')
            layerName = vals[0]
            kern = int(vals[1])
            assert kern == int(allConvLayerNamesParams[layerName]['kernW'])
            assert kern == int(allConvLayerNamesParams[layerName]['kernH'])
            inChans = int(vals[2])
            assert inChans == int(allConvLayerNamesParams[layerName]['inChans'])
            outChans = int(vals[3])
     #       assert outChans == int(allConvLayerNamesParams[layerName]['outChans'])

            srcFullSectNum = int(allConvLayerNamesParams[layerName]['srcFullSectNum'])
            srcReplSectNum = int(allConvLayerNamesParams[layerName]['srcReplSectNum'])
            srcReplUnitNum = int(allConvLayerNamesParams[layerName]['srcReplUnitNum'])
            srcReplUnitWidth = int(allConvLayerNamesParams[layerName]['srcReplUnitWidth'])
            convHalfRateMode = int(allConvLayerNamesParams[layerName]['convHalfRateMode'])

            vals = vals[4:]
            weights = [float(v) for v in vals]

        fname = "%s/fwbqb_bias_%d" % (args['datadir'], fi)
        with open(fname, 'r') as f:
            data = f.read()
            vals = data.strip().split(' ')
            vals = vals[4:]
            bias = [float(v) for v in vals]

        if layer2OffsetMapStr != "":
            layer2OffsetMapStr += ","
        layer2OffsetMapStr += "%s:%d" % (layerName, offset)

        offset += xdnn.v3fillWeightsBiasQuantBlob(blob, offset,
                args['quantizecfg'], weights, args['scaleA'], bias, args['scaleB'],
                kern, kern, inChans, outChans, srcFullSectNum, srcReplSectNum, srcReplUnitNum, srcReplUnitWidth, convHalfRateMode, layerName)

        fi += 1
    xdnn.loadBlobToDdr(blob, size, layer2OffsetMapStr, int(args['PE']))

    return blob

def loadWeightsBiasQuant(args):
    print("Loading weights/bias/quant_params to FPGA...")

    size = 0
    fi = 0
    while True:
        fname = "%s/wbq_size_%d" % (args['datadir'], fi)
        if not os.path.isfile(fname):
            break

        with open(fname, 'r') as f:
            data = f.read()
            vals = data.strip().split(' ')
            vals = [int(v) for v in vals]
            if 'v2WeightsFormat' in args and args['v2WeightsFormat'] == 1:
                size += xdnn.computeWeightsBiasQuantSize(\
                    vals[0], vals[1], vals[2], vals[3], True if args['quantizecfg'] else False)
            else:
                size += xdnn.computeWeightsBiasQuantSize(\
                    vals[0], vals[0], vals[1], vals[2], True if args['quantizecfg'] else False)

        fi += 1

    # print "ANDBG total A size %d" % size
    blob = xdnn.makeWeightsBiasQuantBlob(size)

    fi = 0
    offset = 0
    layer2OffsetMapStr = ""
    while True:
        fname = "%s/fwbqb_%d" % (args['datadir'], fi)
        if 'binaryFormatWeights' in args:
            if args['binaryFormatWeights'] == 1:
                fname = "%s/fwbqbbinary_%d.bin" % (args['datadir'], fi)
        if not os.path.isfile(fname):
            break

        layerName = None
        weights = None
        bias = None
        if 'v2WeightsFormat' in args and args['v2WeightsFormat'] == 1:
            kernWidth = None
            kernHeight = None
        else:
            kern = None
        inChans = None
        outChans = None

        if 'binaryFormatWeights' in args:
            if args['binaryFormatWeights'] == 1:
                if 'v2WeightsFormat' in args and args['v2WeightsFormat'] != 1:
                    print("Weights should be given in the v2WeightsFormat, KernSizex and KernSizey. If they are already\
                        in v2WeightsFormat, then set the argument v2WeightsFormat to 1")
                    sys.exit(1)
                vals = []
                (layerName, kernWidth, kernHeight, inChans, outChans, vals) = xdnn.readWeightsFile(fname)
                weights = [float(v) for v in vals]
            elif args['binaryFormatWeights'] == 0:
                with open(fname, 'r') as f:
                    data = f.read()
                    vals = data.strip().split(' ')
                    layerName = vals[0]
                    if 'v2WeightsFormat' in args and args['v2WeightsFormat'] == 1:
                        kernWidth = int(vals[1])
                        kernHeight = int(vals[2])
                        inChans = int(vals[3])
                        outChans = int(vals[4])
                        vals = vals[5:]
                    else:
                        kern = int(vals[1])
                        inChans = int(vals[2])
                        outChans = int(vals[3])
                        vals = vals[4:]
                    weights = [float(v) for v in vals]
        else:
            with open(fname, 'r') as f:
                data = f.read()
                vals = data.strip().split(' ')
                layerName = vals[0]
                if 'v2WeightsFormat' in args and args['v2WeightsFormat'] == 1:
                    kernWidth = int(vals[1])
                    kernHeight = int(vals[2])
                    inChans = int(vals[3])
                    outChans = int(vals[4])
                    vals = vals[5:]
                else:
                    kern = int(vals[1])
                    inChans = int(vals[2])
                    outChans = int(vals[3])
                    vals = vals[4:]
                weights = [float(v) for v in vals]

        fname = "%s/fwbqb_bias_%d" % (args['datadir'], fi)
        with open(fname, 'r') as f:
            data = f.read()
            vals = data.strip().split(' ')
            if 'v2WeightsFormat' in args:
                if args['v2WeightsFormat'] == 1:
                    vals = vals[5:]
                elif args['v2WeightsFormat'] == 0:
                    vals = vals[4:]
            else:
                vals = vals[4:]
            bias = [float(v) for v in vals]

        if layer2OffsetMapStr != "":
            layer2OffsetMapStr += ","
        layer2OffsetMapStr += "%s:%d" % (layerName, offset)

        # print "ANDBG %d %s %d %d %d" % (offset, layerName, fi, len(weights), len(bias))
        if 'v2WeightsFormat' in args and args['v2WeightsFormat'] == 1:
            offset += xdnn.fillWeightsBiasQuantBlob(blob, offset,
                args['quantizecfg'], weights, args['scaleA'], bias, args['scaleB'],
                kernWidth, kernHeight, inChans, outChans, layerName)
        else:
            offset += xdnn.fillWeightsBiasQuantBlob(blob, offset,
                args['quantizecfg'], weights, args['scaleA'], bias, args['scaleB'],
                kern, kern, inChans, outChans, layerName)

        fi += 1
    xdnn.loadBlobToDdr(blob, size, layer2OffsetMapStr, int(args['PE']))

    return blob

def getNearFileMatchWithPrefix(path, prefix):
    nearMatches = [f for f in os.listdir(path) if f.startswith(prefix)]
    nearMatches.sort()
    if len(nearMatches) > 0:
        return "%s/%s" % (path, nearMatches[0])

    return None

def loadFCWeightsBias(data_dir):
    weight = []
    fname = "%s/fc" % data_dir
    if not os.path.exists(fname):
        nearMatch = getNearFileMatchWithPrefix(data_dir, "fc")
        if nearMatch:
            fname = nearMatch
    if os.path.exists(fname):
        with open(fname, 'r') as f:
            line = f.read()
            vals = line.strip().split(' ')
            weight = [float(v) for v in vals]
    else:
        print(("No FC layers found in %s" % data_dir))
        return (None, None)

    bias = []
    fname = "%s/fc_bias" % data_dir
    if not os.path.exists(fname):
        nearMatch = getNearFileMatchWithPrefix(data_dir, "fc_bias")
        if nearMatch:
            fname = nearMatch
    with open(fname, 'r') as f:
        line = f.read()
        vals = line.strip().split(' ')
        bias = [float(v) for v in vals]

    (weightFpga, fpgaHandle) = xdnn.makeFPGAFloatArray(len(weight))
    (biasFpga, fpgaHandle) = xdnn.makeFPGAFloatArray(len(bias))

    weightFpga[:] = weight
    biasFpga[:] = bias

    return (weightFpga, biasFpga)

# def default_rt_args():
        # parser = argparse.ArgumentParser()

def default_compiler_args():
        parser = argparse.ArgumentParser()
        parser.add_argument("--memory", type=int, default=8)
        parser.add_argument("--generatefile", type=str)
        parser.add_argument("--dsp", type=int, default=28)
        parser.add_argument("--ddr", type=int, default=4096)
        parser.add_argument("--versionjson", type=str, default=None)
        parser.add_argument("--fromtensorflow", default=True)
        parser.add_argument("--weights", default=False)
        parser.add_argument("--strategy", type=str, default="all")
        parser.add_argument("--schedulefile", type=str)
        parser.add_argument("--pngfile", type=str)
        parser.add_argument("--rankdir", type=str, default='BT')
        parser.add_argument("--bytesperpixels", type=int, default=1)
        parser.add_argument("--finalnode", type=str, default=None)
        parser.add_argument("--networkfile", type=str, required=True)
        parser.add_argument("--placeholdershape", type=str, default=None)
        parser.add_argument("--verbose", type=bool, default=False)
        parser.add_argument("--lasttensorbyname", type=str, default=None)
        parser.add_argument("--concatstrategy", type=str, default=None)
        parser.add_argument("--dedicateddsp", type=bool, default=False)

        parser.add_argument("--device", type=str, required=True, choices=['CPU', 'FPGA', 'HWEmu'])
        return parser


import subprocess
import os
from multiprocessing import Pool
from SUPERSTAR.config import SuperStarConfig
from SUPERSTAR.models import PlifProbeScore, PlifProbeGridScoreNew
from rdkit import Chem
from rdkit.Chem import MCS
from rdkit.Chem import AllChem
import math
from IOhandle.models import Protein, Molecule, Target
from PLIFS.models import PlifProbe
from PLIFS.functions import restore_probs
from SUPERSTAR.parse_grid import parse_grid_file, interpolate
from SUPERSTAR.functions import make_json
import time, ast
from django.core.exceptions import ValidationError
import numpy as np
import threading

num_workers = 0
p_dict = {}

def get_centroid(rdmol):
    """Function to return the centroid of a molecule"""
    try:
        for conf in rdmol.GetConformers():
            centroid = Chem.rdMolTransforms.ComputeCentroid(conf)
            x = centroid[0]
            y = centroid[1]
            z = centroid[2]
        return "%s %s %s" % (x, y, z)
    except AttributeError:
        return False

def icm_add_hyds(pdb_code):
    """ """
    out_file = 'W:/Informatics/Pharmacophore/anthony/DPhil/CODE/CHOC/src/WebApp/SUPERSTAR/PROTEINS/' + pdb_code + "_OUT.pdb"
    pdb_file = 'W:/Informatics/Pharmacophore/anthony/DPhil/CODE/CHOC/src/WebApp/SUPERSTAR/PROTEINS/' + pdb_code + ".pdb"
    out_icm = """call _macro
# Load the file_in
openFile '""" + pdb_file + """'
# Convert
convertObject a_1.* yes yes yes yes yes yes no
#Save it
my_s = '""" + out_file + """'
write pdb a_1.* my_s
q
    """
    out_f = open("out_icm", "w")
    out_f.write(out_icm)
    out_f.close()
    #subprocess.call(["C:\Program Files (x86)\Molsoft LLC\ICM-Pro\icmconsole.exe", "out_icm"])
    return out_file

def run_worker(jobname, ss_c, probe, centroid, mol, rd_conf, h_path, my_probes, grid_space):
    """Function to run each thread as different thread"""
    # NOW DEFINE THE JOBNAME
    global num_workers
    global p_dict
    num_workers += 1
    print jobname, "WORKER", num_workers
    ins_file = '''JOBNAME ''' + jobname + '''
PROBENAME ''' + ss_c.probe_dict[probe] + '''
GRID_SPACING ''' + str(grid_space) + '''
CALC_CONTOUR_SURFACE 0
RUN_TYPE SUPERSTAR
DATA_SOURCE CSD
SHELL_VALUES 0.5 0
SIGCHECK_METHOD POISSON SHELL
PROPENSITY_CORRECTION LOGP DEFAULT DEFAULT
MIN_CAVITY_VOLUME 1
CAVITY_RADIUS 3
DATA_TYPE ISOSTAR
MOLECULE_FILE ''' + h_path + '''
CAVITY_DETECTION 1
MIN_PSP 4
MIN_PROPENSITY 0.0
SAVE_CAVITY MESH
USE_TORSIONS 1
MAP_FORMAT INSIGHT_ASCII_GRID

  '''
    # Now check if there is a centroid
    if centroid:
        ins_file = ins_file + "CAVITY_ORIGIN " + str(centroid) + "\n"
    # Now define the ins file
    ins_file_path = os.path.join(ss_c.ins_path, jobname + ".ins")
    input_file = open(ins_file_path, "w")
    input_file.write(ins_file)
    input_file.close()
    s_path = r"C:\Program Files (x86)\CCDC\SuperStar 2.1.2\bin\superstar_app.exe"
    cmd = [s_path, ins_file_path]
    # Now run it
    in_path = os.path.join(ss_c.out_path, jobname + ".grd")
    print "SPAWNING SUPERSTAR"
    if not os.path.isfile(in_path):
        subprocess.Popen(cmd, shell=True)
    # Wait for the files to write
    while 1 == 1:
        if not os.path.isfile(in_path):
            time.sleep(1)
            continue
        break
    # Now get the grid back
    print "PARSING GRID"
    grid_space, grid_dict = parse_grid_file(in_path)
    print "GRID PARSED"
    # Now read in the grid
    # Now go through the relevant PLIF probes
    # Now loop through these and score them
    for p in my_probes:
        # Get the atom ids
        propensities = []
        atom_ids = ast.literal_eval(p.atom_ids)
        for atom_id in atom_ids:
            # Get the atom
            my_atom = rd_conf.GetAtomPosition(atom_id)
            # Now loop through the atom ids
            propensities.append(interpolate(grid_space, grid_dict, my_atom.x, my_atom.y, my_atom.z))
        # Now get the geometric mean of this
        propensity = np.exp(np.sum(np.log(propensities))) / len(atom_ids)
        p_dict[p.pk] = propensity
    num_workers -= 1
    return None

class myThread (threading.Thread):
    """Class to start new threads"""
    def __init__(self, threadID, jobname, ss_c, probe, centroid, mol, rd_conf, h_path, my_probes, grid_space):
        threading.Thread.__init__(self)
        self.threadID = threadID
        self.name = jobname
        #self.counter = counter
        self.jobname = jobname
        self.ss_c = ss_c
        self.probe = probe
        self.centroid = centroid
        self.mol = mol
        self.rd_conf = rd_conf
        self.h_path = h_path
        self.my_probes = my_probes
        self.grid_space = grid_space

    def run(self):
        print "Starting " + self.name
        run_worker(self.jobname, self.ss_c, self.probe, self.centroid, self.mol, self.rd_conf, self.h_path, self.my_probes, self.grid_space)

def make_ins_file(target_id=2, grid_space=0.5):
    """Function to take a target and do analysis of all pharma points bound to this target
    in xtal structures"""
    global p_dict
    global num_workers
    ss_c = SuperStarConfig()
    target = Target.objects.get(pk=target_id)
    my_mols = Molecule.objects.filter(prot_id__target_id=target_id).exclude(prot_id__code__contains=target.title)
    # Now change to this directoy
    os.chdir(ss_c.out_path)
    # Loop through the appropriate proteins
    for mol in my_mols:
        # Define the files to clean up
        rd_mol = Chem.MolFromMolBlock(str(mol.sdf_info))
        rd_conf = rd_mol.GetConformer()
        file_list = []
        # Define the pdb_file
        pdb = mol.prot_id.code
        out_path = os.path.join(ss_c.base_path, pdb + ".pdb")
        out_f = open(out_path, "w")
        pdb_mol = Chem.MolFromPDBBlock(str(mol.prot_id.pdb_info))
        out_f.write(Chem.MolToPDBBlock(pdb_mol))
        out_f.close()
        out_file = icm_add_hyds(pdb)
        # Now use ICM to add Hydrogens to this...
        h_path = os.path.join(ss_c.base_path, pdb + "_OUT.pdb")
        centroid = get_centroid(rd_mol)
        my_keys = ss_c.probe_dict.keys()
        counter = -1
        while 1 == 1:
            if num_workers >= 6:
                continue
            elif counter < len(my_keys)-1:
                counter += 1
                probe = my_keys[counter]
                jobname = pdb + "_" + probe + "_" + str(grid_space).replace(".", "_")
                my_probes = PlifProbe.objects.filter(mol_id=mol, type__in=ss_c.type_dict[probe])
                len(my_probes)
                my_t = myThread(counter, jobname, ss_c, probe, centroid, mol, rd_conf, h_path, my_probes, grid_space)
                my_t.start()
            if num_workers == 0 and counter == len(my_keys) -1:
                break
    # Now add the scores - with this multithreaded EG
    for p in p_dict:
        p_score = PlifProbeGridScoreNew()
        plif_p = PlifProbe.objects.get(pk=p)
        p_score.plif_probe = plif_p
        p_score.score = p_dict[p]
        p_score.grid_space = grid_space
        try:
            p_score.validate_unique()
            p_score.save()
        except ValidationError:
            p_score = PlifProbeGridScoreNew.objects.get(plif_probe=plif_p, grid_space=grid_space)
            p_score.score = p_dict[p]
            p_score.save()

def run_anal(target_id=2, grid_spaces=[0.1, 0.2, 0.3, 0.4, 0.5,
                                       0.6, 0.7, 0.8, 0.9, 1.0]):
    """Function to run the analysis on multiple grid spacings"""
    global p_dict
    # Loop over the options
    for grid_s in grid_spaces:
        # Run the function
        # Re-zero the dictionary
        p_dict = {}
        make_ins_file(target_id, grid_space=grid_s)

def run_targs():
    """Function to run the SUPERSTAR analysis for all targets"""
    targs = Target.objects.exclude(title="DUMMY")
    for targ in targs:
        # Now run the analysis for this target
        restore_probs(targ.pk)
        run_anal(targ.pk, grid_spaces=[0.5])
        make_json(targ, grid_space=0.5)


"""ssd.py
This module implements the TrtSSD class.
"""
import ctypes
import numpy as np
import cv2
import tensorrt as trt
import pycuda.driver as cuda
import openpifpaf
import torch
import PIL
import sys
import time
import cv2 as cv

class PoseEstimator():

    def _load_engine(self):
        TRTbinPath = self.trt_bin_path
        with open(TRTbinPath, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
            return runtime.deserialize_cuda_engine(f.read())

    #def _create_context(self):
    def _allocate_buffers(self):
        summ = 0
        for i in range(self.engine.num_bindings):
            binding = self.engine[i]
            size = trt.volume(self.engine.get_binding_shape(binding)) * \
                   self.engine.max_batch_size
            summ = summ + size
            host_mem = cuda.pagelocked_empty(size, np.float32)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(cuda_mem))
            if self.engine.binding_is_input(binding):
                self.host_inputs.append(host_mem)
                self.cuda_inputs.append(cuda_mem)
            else:
                self.host_outputs.append(host_mem)
                self.cuda_outputs.append(cuda_mem)
        return

    def __init__(self, trt_bin_path, model_input_size):
        """Initialize TensorRT plugins, engine and conetxt."""
        self.trt_bin_path = trt_bin_path
        self.model = ''
        self.trt_logger = trt.Logger(trt.Logger.INFO)
        self.model_input_size = model_input_size
        self.host_inputs = []
        self.cuda_inputs = []
        self.host_outputs = []
        self.cuda_outputs = []
        self.bindings = []
        self._init_cuda_stuff()

    def _init_cuda_stuff(self):
        cuda.init()
        self.device = cuda.Device(0)  # enter your Gpu id here
        self.cuda_context = self.device.make_context()
        self.engine = self._load_engine()
        self._allocate_buffers()
        self.engine_context = self.engine.create_execution_context()
        self.stream = cuda.Stream()  # create a CUDA stream to run inference

    def __del__(self):
        """ Free CUDA memories. """
        for mem  in self.cuda_inputs:
            mem.free()
        for mem in self.cuda_outputs:
            mem.free

        del self.stream
        del self.cuda_outputs
        del self.cuda_inputs
        self.cuda_context.pop()
        del self.cuda_context
        del self.engine_context
        del self.engine
        del self.bindings
        del self.host_inputs
        del self.host_outputs

    def inference(self, img):
        """Estimate human poses in the input image."""
        self._init_cuda_stuff()

        img = cv.resize(img, self.model_input_size)
        pil_im = PIL.Image.fromarray(img)
        preprocess = None
        data = openpifpaf.datasets.PilImageList([pil_im], preprocess=preprocess)

        loader = torch.utils.data.DataLoader(
            data, batch_size=1, shuffle=False,
            pin_memory=True,
            collate_fn=openpifpaf.datasets.collate_images_anns_meta)
        for images_batch, _ , __ in loader:
            np_img = images_batch.numpy() # np_img.size = (0,3,641,641)

        self.host_inputs[0] = np.ravel(np.zeros_like(np_img)) #np.ravel(np.zeros_like(np_img))

        start_time = time.time()

        np.copyto(self.host_inputs[0], np.ravel(np_img))
        cuda.memcpy_htod_async(
            self.cuda_inputs[0], self.host_inputs[0], self.stream)
        self.stream.synchronize()

        self.engine_context.execute_async(
            batch_size=1,
            bindings=self.bindings,
            stream_handle=self.stream.handle)

        pif_names=['pif_c', 'pif_r', 'pif_b', 'pif_s']
        paf_names=['paf_c','paf_r1','paf_r2','paf_b1','paf_b2']
        other_names=['616','646','648','637','644']
        pif=[None]*4
        paf=[None]*5
        other=[None]*5
        for i in range(1, self.engine.num_bindings): #int(len(self.host_outputs)/2) +1):
            cuda.memcpy_dtoh_async(
                self.host_outputs[i - 1], self.cuda_outputs[i - 1], self.stream)

        self.stream.synchronize()
        print("--- %s seconds ---" % (time.time() - start_time))

        for i in range(1, self.engine.num_bindings): #int(len(self.host_outputs)/2) +1):
            shape = self.engine.get_binding_shape(i)
            name = self.engine.get_binding_name(i)
            total_shape = np.prod(shape)
            output = self.host_outputs[i - 1][0: total_shape]
            output = np.reshape(output, tuple(shape))
            if name in pif_names:
                index_n = pif_names.index(name)
                pif[index_n] = torch.from_numpy(output)
            elif name in paf_names:
                index_n = paf_names.index(name)
                paf[index_n] = torch.from_numpy(output)
            elif name in other_names:
                index_n = other_names.index(name)
                other[index_n] = torch.from_numpy(output)

        final_output = [pif, paf, other]
        self.cuda_context.pop()

        return final_output


# importing above defined libraries to
# implement the functionalities
from PIL import ImageGrab, ImageOps
import pyautogui
import time
import numpy as np

class cordinates():

	# coordinates of replay button to start the game
	replaybutton =(360, 214)
	# this coordinates represent the top-right coordinates
	# that will be used to define the front box
	dinasaur = (149, 239 )

def restartGame():

	# using pyautogui library, we are clicking on the
	# replay button without any user interaction
	pyautogui.click(cordinates.replaybutton)

	# we will keep our Bot always down that
	# will prevent him to get hit by bird
	pyautogui.keyDown('down')

def press_space():

	# releasing the Down Key
	pyautogui.keyUp('down')

	# pressing Space to overcome Bush
	pyautogui.keyDown('space')

	# so that Space Key will be recognized easily
	time.sleep(0.05)

	# printing the "Jump" statement on the
	# terminal to see the current output
	print("jump")
	time.sleep(0.10)

	# releasing the Space Key
	pyautogui.keyUp('space')

	# again pressing the Down Key to keep my Bot always down
	pyautogui.keyDown('down')

def imageGrab():
	# defining the coordinates of box in front of dinosaur
	box = (cordinates.dinasaur[0]+30, cordinates.dinasaur[1],
		cordinates.dinasaur[0]+120, cordinates.dinasaur[1]+2)

	# grabbing all the pixels values in form of RGB tupples
	image = ImageGrab.grab(box)

	# converting RGB to Grayscale to
	# make processing easy and result faster
	grayImage = ImageOps.grayscale(image)

	# using numpy to get sum of all grayscale pixels
	a = np.array(grayImage.getcolors())

	# returning the sum
	print(a.sum())
	return a.sum()

# function to restart the game
restartGame()
while True:
	# 435 is the sum of white pixels values of box.
	# You may get different value is you are taking bigger
	# or smaller box than the box taken in this article.
	# if value returned by "imageGrab" function is not equal to 435,
	# it means either bird or bush is coming towards dinosaur
	if(imageGrab()!= 435):
		press_space()
		# time to recognize the operation performed by above function
		time.sleep(0.1)


"""tau IDs"""

import numpy as np
from fnal_column_analysis_tools.util import awkward
from fnal_column_analysis_tools.analysis_objects import JaggedCandidateArray

# descriptions from https://twiki.cern.ch/twiki/bin/view/CMSPublic/SWGuidePFTauID#Discriminators
kByDecayModeFinding                        = 1 << 0   # You will always want to use this (see AN-10-82)

kByVLooseIsolation                         = 1 << 1   # isolation cone of 0.3, no PF Charged Candidates with pT > 1.5 GeV/c and no PF Gamma candidates with ET > 2.0 GeV
kByLooseIsolation                          = 1 << 2   # (description N/A)
kByMediumIsolation                         = 1 << 3   # (description N/A)
kByTightIsolation                          = 1 << 4   # (description N/A)
kByVLooseIsolationDBSumPtCorr              = 1 << 5   # (description N/A)
kByLooseIsolationDBSumPtCorr               = 1 << 6   # (description N/A)
kByMediumIsolationDBSumPtCorr              = 1 << 7   # (description N/A)
kByTightIsolationDBSumPtCorr               = 1 << 8   # (description N/A)
kByVLooseCombinedIsolationDBSumPtCorr      = 1 << 9   # isolation cone of 0.3, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 3 GeV
kByLooseCombinedIsolationDBSumPtCorr       = 1 << 10  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 2 GeV
kByMediumCombinedIsolationDBSumPtCorr      = 1 << 11  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 1 GeV
kByTightCombinedIsolationDBSumPtCorr       = 1 << 12  # isolation cone of 0.5, Delta Beta corrected sum pT of PF charged and PF gamma isolation candidates (pT > 0.5 GeV) less than 0.8 GeV
kByLooseCombinedIsolationDBSumPtCorr3Hits  = 1 << 13  # same as ByLooseCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates
kByMediumCombinedIsolationDBSumPtCorr3Hits = 1 << 14  # same as ByMediumCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates
kByTightCombinedIsolationDBSumPtCorr3Hits  = 1 << 15  # same as ByTightCombinedIsolationDBSumPtCorr but requiring 3 hits (instead of 8) on track of isolation candidates
kByLooseIsolationMVA                       = 1 << 16  # BDT based selection using isolation in rings around tau direction and shower shape variables
kByMediumIsolationMVA                      = 1 << 17  #  "
kByTightIsolationMVA                       = 1 << 18  #  "
kByLooseIsolationMVA2                      = 1 << 19  # same as ByLooseIsolationMVA with new training and improved performance
kByMediumIsolationMVA2                     = 1 << 20  # same as ByMediumIsolationMVA with new training and improved performance
kByTightIsolationMVA2                      = 1 << 21  # same as ByTightIsolationMVA with new training and improved performance

kByLooseElectronRejection                  = 1 << 22  # electron pion MVA discriminator < 0.6
kByMediumElectronRejection                 = 1 << 23  # electron pion MVA discriminator < -0.1 and not 1.4442 < |eta| < 1.566
kByTightElectronRejection                  = 1 << 24  # electron pion MVA discriminator < -0.1 and not 1.4442 < |eta| < 1.566 and Brem pattern cuts (see AN-10-387)
kByMVA3LooseElectronRejection              = 1 << 25  # anti-electron MVA discriminator with improved training
kByMVA3MediumElectronRejection             = 1 << 26  #  "
kByMVA3TightElectronRejection              = 1 << 27  #  "
kByMVA3VTightElectronRejection             = 1 << 28  #  "

kByLooseMuonRejection                      = 1 << 29  # Tau Lead Track not matched to chamber hits
kByMediumMuonRejection                     = 1 << 30  # Tau Lead Track not matched to global/tracker muon
kByTightMuonRejection                      = 1 << 31  # Tau Lead Track not matched to global/tracker muon and large enough energy deposit in ECAL+HCAL exceeding 0.2 times Lead Track momentum
kByLooseMuonRejection2                     = 1 << 32  # Same as LooseMuonRejection
kByMediumMuonRejection2                    = 1 << 33  # LooseMuonRejection2 && no DT, CSC or RPC Hits in last 2 Stations
kByTightMuonRejection2                     = 1 << 34  # MediumMuonRejection2 && large enough energy deposit in ECAL+HCAL in 1 prong + 0 strip decay mode (SUM(ECAL+HCAL)>0.2*pT)
kByLooseMuonRejection3                     = 1 << 35  # Same as LooseMuonRejection
kByTightMuonRejection3                     = 1 << 36  # MediumMuonRejection2 && large enough energy deposit in ECAL+HCAL in 1 prong + 0 strip decay mode (SUM(ECAL+HCAL)>0.2*pT)

def passTauSel(tau):
    return ( ((tau.hpsDisc & kByDecayModeFinding) > 0) & (tau.rawIso3Hits < 5) )


__author__ = "tindht@vietmoney.vn"
__copyright__ = "Copyright 2021, VietMoney Face Anti-spoofing"

from typing import Any, Sequence, Tuple

import cv2
import nptyping as npt
import numpy as np
import torch
from torch.nn.functional import softmax

from library.models.mini_fasnet import MiniFASNetV1SE, MiniFASNetV2
from library.util import scale_box, remove_prefix
from library.util.transform import Transform, resize

__all__ = ['MulFasNet', 'SpoofingDetector']
_CPU_DEVICE = "cpu"

class MulFasNet:
    """
    Ref: https://github.com/minivision-ai/Silent-Face-Anti-Spoofing

    The face anti spoofing modify from Silent-Face-Anti-Spoofing with optimize data flow and minimal coding.
    `MiniFASNetV1SE` and `MiniFASNetV2` have been compiled into single model.
    """

    def __init__(self, model_path: str, device=_CPU_DEVICE, input_size=(80, 80)):
        """
        Parameters:
            model_path: pretrained model
            device: device model loaded in.
            input_size: model input size.
        """
        model_load = torch.load(model_path, map_location=device)
        kernel_size = ((input_size[0] + 15) // 16, (input_size[1] + 15) // 16)
        self.models = [MiniFASNetV1SE(conv6_kernel=kernel_size), MiniFASNetV2(conv6_kernel=kernel_size)]

        # Load model
        for idx, model in enumerate(self.models):
            model.load_state_dict(remove_prefix(model_load[model.__class__.__name__], 'module.'), strict=False)
            model.to(device)
            model.eval()

        self.max_scale = [4, 2.7]
        self.device = device

    def forward(self, faces_scale: Sequence[torch.Tensor]) -> npt.NDArray[(Any, 3), npt.Float]:
        """
        Predict face spoof.

        Parameters
        ----------
            faces_scale: face's image scaled.

        Returns
        -------
            Score of predict [2D_SPOOF | REAL | 3D_SPOOF]
        """
        predict = torch.zeros(faces_scale[0].size(0), 3, dtype=torch.float64, device=self.device)
        with torch.no_grad():
            for model, images in zip(self.models, faces_scale):
                images = images.to(self.device)
                result = model(images)
                result = softmax(result, dim=1)
                predict += result

        if self.device != _CPU_DEVICE:
            predict = predict.detach().cpu()
        return predict.detach().numpy()

    __call__ = forward

class SpoofingDetector:
    """
    Predict faces in image that is real or not.

    Parameters:
        model_path: pretrained model
        device: device model loaded in. (Default: cpu)
        face_size: model face input size.
    """

    def __init__(self, model_path: str, device: str = _CPU_DEVICE, face_size=(80, 80)):
        self.model = MulFasNet(model_path, device=device, input_size=face_size)
        self.device = self.model.device
        self.face_size = face_size

        self.transform = Transform([
            lambda x: cv2.cvtColor(x, cv2.COLOR_RGB2BGR),
            resize(*face_size),
            lambda x: torch.from_numpy(x.transpose((2, 0, 1))).float()
        ])

    def predict(self, boxes: Sequence[Sequence[int]],
                image: npt.NDArray[npt.UInt8]) -> Sequence[Tuple[bool, float]]:
        """
        Post-process predict from model. Calculate average score and label face.

        Parameters
        ----------
            boxes: Face's boxes
            image: image source

        Returns
        -------
            Label and score of faces in images.
            [True|False] == [REAL|FAKE]
        """
        faces_scale = list()
        for scale in self.model.max_scale:
            face_tensor = torch.zeros(len(boxes), 3, *self.face_size, dtype=torch.float32)
            for idx, box in enumerate(boxes):
                height, width = image.shape[:2]
                box_scaled = scale_box(width, height, box, scale)
                face_img = image[box_scaled[1]:box_scaled[3], box_scaled[0]:box_scaled[2]]
                face_tensor[idx] = self.transform(face_img)
            faces_scale.append(face_tensor)

        predict = self.model(faces_scale)
        labels = np.argmax(predict, axis=1)
        return [(label == 1, predict[idx][label] / 2.) for idx, label in enumerate(labels)]

    __call__ = predict