# Copyright 2020 Graphcore Ltd.
import argparse
import os

import numpy as np
import tensorflow.compat.v1 as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras import layers

from tensorflow.python import ipu
from tensorflow.compiler.plugin.poplar.ops import gen_ipu_ops

try:
    from gcprofile import save_tf_report
    gcprofile_present = True
except ImportError:
    gcprofile_present = False

def parse_args():
    # Handle command line arguments
    pipeline_schedule_options = [p.name for p in ipu.pipelining_ops.PipelineSchedule]

    parser = argparse.ArgumentParser()
    parser.add_argument("--batch-size", type=int, default=32,
                        help="The batch size.")
    parser.add_argument("--repeat-count", type=int, default=100,
                        help="The number of times the pipeline will be executed for each step.")
    parser.add_argument("--epochs", type=float, default=50,
                        help="Total number of epochs to train for.")
    parser.add_argument("--learning-rate", type=float, default=0.01,
                        help="The learning rate used with stochastic gradient descent.")
    parser.add_argument('--pipeline-depth', type=int, default=16,
                        help="Depth of pipeline to use.")
    parser.add_argument('--pipeline-schedule', type=str, default="Grouped",
                        choices=pipeline_schedule_options,
                        help="Pipelining scheduler. In the 'Grouped' schedule the forward passes"
                        " are grouped together, and the backward passes are grouped together. "
                        "With 'Interleaved' the forward and backward passes are interleaved. "
                        "'Sequential' mimics a non-pipelined execution.")
    parser.add_argument("--synthetic-data", action='store_true',
                        help="Use synthetic data instead of real images.")
    parser.add_argument("--profile", action='store_true',
                        help="Enable profiling for gc-profile.")
    args = parser.parse_args()
    return args

def create_dataset(args):
    # Prepare a tf dataset with mnist data
    train_data, _ = mnist.load_data()

    def normalise(x, y):
        return x.astype("float32") / 255.0, y.astype("int32")

    x_train, y_train = normalise(*train_data)

    def generator():
        return zip(x_train, y_train)

    types = (x_train.dtype, y_train.dtype)
    shapes = (x_train.shape[1:], y_train.shape[1:])

    n_examples = len(x_train)
    dataset = tf.data.Dataset.from_generator(generator, types, shapes)
    # Caching and prefetching are important to prevent the host data
    # feed from being the bottleneck for throughput
    dataset = dataset.batch(args.batch_size, drop_remainder=True)
    dataset = dataset.shuffle(n_examples)
    dataset = dataset.cache()
    dataset = dataset.repeat()
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return n_examples, dataset

# The following is a schematic representation of the model defined in this example,
# which also shows how it is split across two IPUs:
# ------------------------------------ Model Definition ------------------------------------
#  <----------------------- ipu0 -----------------------> <------------- ipu1 ------------->
#
# inputs --|-- Dense --|-- Relu --|-- Dense --|-- Relu --|-- Dense--|-- SoftmaxCE --|-- Loss
#                 w0 --|                 w1 --|                w2 --|
#                 b0 --|                 b1 --|                b2 --|

def stage1(lr, images, labels):
    # Stage 1 of the pipeline. Will be placed on the first IPU
    partial = layers.Flatten()(images)
    partial = layers.Dense(256, activation=tf.nn.relu)(partial)
    partial = layers.Dense(128, activation=tf.nn.relu)(partial)
    return lr, partial, labels

def stage2(lr, partial, labels):
    # Stage 2 of the pipeline. Will be placed on the second IPU
    logits = layers.Dense(10)(partial)
    cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=labels, logits=logits)
    loss = tf.reduce_mean(cross_entropy)
    return lr, loss

def optimizer_function(lr, loss):
    # Optimizer function used by the pipeline to automatically set up
    # the gradient accumulation and weight update steps
    optimizer = tf.train.GradientDescentOptimizer(lr)
    return ipu.pipelining_ops.OptimizerFunctionOutput(optimizer, loss)

def model(lr):
    # Defines a pipelined model which is split accross two stages
    with tf.variable_scope("FCModel", use_resource=True):
        pipeline_op = ipu.pipelining_ops.pipeline(
            computational_stages=[stage1, stage2],
            pipeline_depth=args.pipeline_depth,
            repeat_count=args.repeat_count,
            inputs=[lr],
            infeed_queue=infeed_queue,
            outfeed_queue=outfeed_queue,
            optimizer_function=optimizer_function,
            pipeline_schedule=next(p for p in ipu.pipelining_ops.PipelineSchedule
                                   if args.pipeline_schedule == p.name),
            outfeed_loss=True,
            name="Pipeline")
    return pipeline_op

if __name__ == "__main__":
    args = parse_args()

    if args.synthetic_data:
        if "TF_POPLAR_FLAGS" in os.environ:
            os.environ["TF_POPLAR_FLAGS"] += " --use_synthetic_data"
        else:
            os.environ["TF_POPLAR_FLAGS"] = "--use_synthetic_data"

    n_examples, dataset = create_dataset(args)

    # Create the data queues from/to IPU
    infeed_queue = ipu.ipu_infeed_queue.IPUInfeedQueue(dataset, "infeed")
    outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue("outfeed")

    # With batch size BS, pipeline depth PD and repeat count RPT,
    # at every step n = (BS * PD * RPT) examples are used.
    # So in order to evaluate at least N total examples, do ceil(N / n) steps
    num_train_examples = int(args.epochs * n_examples)
    examples_per_step = args.batch_size * args.pipeline_depth * args.repeat_count
    steps = ((num_train_examples - 1) // examples_per_step) + 1

    with tf.device('cpu'):
        lr = tf.placeholder(np.float32, [])
        report = gen_ipu_ops.ipu_event_trace()

    with ipu.scopes.ipu_scope("/device:IPU:0"):
        compiled_model = ipu.ipu_compiler.compile(model, inputs=[lr])

    outfeed_op = outfeed_queue.dequeue()

    ipu.utils.move_variable_initialization_to_cpu()
    init_op = tf.global_variables_initializer()

    # Configure the IPU
    # With pipelining, IPU-level profiling is needed to correctly visualise the execution trace.
    # For pipelined models either SNAKE or HOOF IPU selection orders are advised;
    # the latter works best when the first and last stage live on the same IPU.
    # For more info, check ipu.utils.py file or the TensorFlow document:
    # https://www.graphcore.ai/docs/targeting-the-ipu-from-tensorflow#tensorflow.python.ipu.utils.SelectionOrder.
    cfg = ipu.utils.create_ipu_config(
        profiling=args.profile,
        profile_execution=ipu.utils.ExecutionProfileType.IPU_PROFILE if args.profile else False,
        selection_order=ipu.utils.SelectionOrder.SNAKE)
    # Auto select as many IPUs as we want to pipeline across
    cfg = ipu.utils.auto_select_ipus(cfg, 2)
    ipu.utils.configure_ipu_system(cfg)

    with tf.Session() as sess:
        # Initialize
        sess.run(init_op)
        sess.run(infeed_queue.initializer)
        # Run
        for step in range(steps):
            sess.run(compiled_model, {lr: args.learning_rate})
            if args.profile and gcprofile_present:
                raw_reports = sess.run(report)
                save_tf_report(raw_reports)
                break
            # Read the outfeed for the training losses
            losses = sess.run(outfeed_op)
            epoch = float(examples_per_step * step / n_examples)
            print("Epoch {:.1f}, Mean loss: {:.3f}".format(
                epoch, np.mean(losses)))


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

# # Self-Driving Car Engineer Nanodegree
#
# ## Deep Learning
#
# ## Project: Build a Traffic Sign Recognition Classifier
#
# In this notebook, a template is provided for you to implement your functionality in stages, which is required to successfully complete this project. If additional code is required that cannot be included in the notebook, be sure that the Python code is successfully imported and included in your submission if necessary.
#
# > **Note**: Once you have completed all of the code implementations, you need to finalize your work by exporting the iPython Notebook as an HTML document. Before exporting the notebook to html, all of the code cells need to have been run so that reviewers can see the final implementation and output. You can then export the notebook by using the menu above and navigating to  \n",
#     "**File -> Download as -> HTML (.html)**. Include the finished document along with this notebook as your submission.
#
# In addition to implementing code, there is a writeup to complete. The writeup should be completed in a separate file, which can be either a markdown file or a pdf document. There is a [write up template](https://github.com/udacity/CarND-Traffic-Sign-Classifier-Project/blob/master/writeup_template.md) that can be used to guide the writing process. Completing the code template and writeup template will cover all of the [rubric points](https://review.udacity.com/#!/rubrics/481/view) for this project.
#
# The [rubric](https://review.udacity.com/#!/rubrics/481/view) contains "Stand Out Suggestions" for enhancing the project beyond the minimum requirements. The stand out suggestions are optional. If you decide to pursue the "stand out suggestions", you can include the code in this Ipython notebook and also discuss the results in the writeup file.
#
#
# >**Note:** Code and Markdown cells can be executed using the **Shift + Enter** keyboard shortcut. In addition, Markdown cells can be edited by typically double-clicking the cell to enter edit mode.

# ---
# ## Step 0: Load The Data

# In[1]:

# Load pickled data
import pickle
import numpy as np
import pandas as pd
import csv

# TODO: Fill this in based on where you saved the training and testing data

training_file = '../data/train.p'
validation_file= '../data/valid.p'
testing_file = '../data/test.p'
sign_names_file = '../data/signnames.csv'

with open(training_file, mode='rb') as f:
    train = pickle.load(f)
with open(validation_file, mode='rb') as f:
    valid = pickle.load(f)
with open(testing_file, mode='rb') as f:
    test = pickle.load(f)

X_train, y_train = train['features'], train['labels']
X_valid, y_valid = valid['features'], valid['labels']
X_test, y_test = test['features'], test['labels']

# ---
#
# ## Step 1: Dataset Summary & Exploration
#
# The pickled data is a dictionary with 4 key/value pairs:
#
# - `'features'` is a 4D array containing raw pixel data of the traffic sign images, (num examples, width, height, channels).
# - `'labels'` is a 1D array containing the label/class id of the traffic sign. The file `signnames.csv` contains id -> name mappings for each id.
# - `'sizes'` is a list containing tuples, (width, height) representing the original width and height the image.
# - `'coords'` is a list containing tuples, (x1, y1, x2, y2) representing coordinates of a bounding box around the sign in the image. **THESE COORDINATES ASSUME THE ORIGINAL IMAGE. THE PICKLED DATA CONTAINS RESIZED VERSIONS (32 by 32) OF THESE IMAGES**
#
# Complete the basic data summary below. Use python, numpy and/or pandas methods to calculate the data summary rather than hard coding the results. For example, the [pandas shape method](http://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.shape.html) might be useful for calculating some of the summary results.

# ### Provide a Basic Summary of the Data Set Using Python, Numpy and/or Pandas

# In[2]:

### Replace each question mark with the appropriate value.
### Use python, pandas or numpy methods rather than hard coding the results

# TODO: Number of training examples
n_train = pd.DataFrame(y_train).shape[0]

# TODO: Number of validation examples
n_validation = pd.DataFrame(y_valid).shape[0]

# TODO: Number of testing examples.
n_test = pd.DataFrame(y_test).shape[0]

# TODO: What's the shape of an traffic sign image?
image_shape = X_train[0].shape

# TODO: How many unique classes/labels there are in the dataset.
line_count = 0
sign_name_database = []
final_sign_data = []
fo = open(sign_names_file)
with fo as csv_file:
    csv_reader = csv.reader(csv_file, delimiter=',')
    for row in csv_reader:
        if line_count > 0:
            sign_name_database.append(row[1])
            final_sign_data.append([str(line_count - 1), row[1]])
        line_count += 1

n_classes = line_count -1

print("Number of training examples =", n_train)
print("Number of testing examples =", n_test)
print("Image data shape =", image_shape)
print("Number of classes =", n_classes)
fo.close()

# ### Include an exploratory visualization of the dataset

# Visualize the German Traffic Signs Dataset using the pickled file(s). This is open ended, suggestions include: plotting traffic sign images, plotting the count of each sign, etc.
#
# The [Matplotlib](http://matplotlib.org/) [examples](http://matplotlib.org/examples/index.html) and [gallery](http://matplotlib.org/gallery.html) pages are a great resource for doing visualizations in Python.
#
# **NOTE:** It's recommended you start with something simple first. If you wish to do more, come back to it after you've completed the rest of the sections. It can be interesting to look at the distribution of classes in the training, validation and test set. Is the distribution the same? Are there more examples of some classes than others?

# In[6]:

import warnings
warnings.filterwarnings("ignore")
import random
### Data exploration visualization code goes here.
### Feel free to use as many code cells as needed.
import matplotlib.pyplot as plt
from matplotlib import colors
from matplotlib.font_manager import FontProperties
# Visualizations will be shown in the notebook.
get_ipython().run_line_magic('matplotlib', 'inline')

index = random.randint(0, len(X_train))
image = X_train[index].squeeze()

plt.figure(figsize=(1,1))
plt.title(str(y_train[index]) +' = ' +  sign_name_database[y_train[index]])
plt.imshow(image)
plt.show()

# Distributions

fig, axs = plt.subplots(1, 2, tight_layout=True, figsize=(15, 6))
fig.subplots_adjust(wspace=50)
axs[0].set_xlabel('Class')
axs[0].set_ylabel('Frequency')
# the histogram of the data
n, bins, patches = axs[0].hist(y_train, n_classes, normed=0, label='Training set')
# the histogram of the data
n, bins, patches = axs[0].hist(y_test, n_classes, normed=0, label='Test set')
n, bins, patches = axs[0].hist(y_valid, n_classes, normed=0, label='Validation set')

axs[0].legend()
axs[0].set_title('Distribution of classes')

table_subplot = axs[1]

table = table_subplot.table(cellText=final_sign_data, colLabels=['Class', 'Description'], loc='center', cellLoc='center', colColours=['#F3CC32', '#2769BD'])
table.auto_set_font_size(False)
table.set_fontsize(10)
table.auto_set_column_width((-1, 0, 1, 2, 3))
table.auto_set_column_width(True)

for (row, col), cell in table.get_celld().items():
    if (row == 0):
        cell.set_text_props(fontproperties=FontProperties(weight='bold', size=12))
plt.axis('off')
plt.show()

# ----
#
# ## Step 2: Design and Test a Model Architecture
#
# Design and implement a deep learning model that learns to recognize traffic signs. Train and test your model on the [German Traffic Sign Dataset](http://benchmark.ini.rub.de/?section=gtsrb&subsection=dataset).
#
# The LeNet-5 implementation shown in the [classroom](https://classroom.udacity.com/nanodegrees/nd013/parts/fbf77062-5703-404e-b60c-95b78b2f3f9e/modules/6df7ae49-c61c-4bb2-a23e-6527e69209ec/lessons/601ae704-1035-4287-8b11-e2c2716217ad/concepts/d4aca031-508f-4e0b-b493-e7b706120f81) at the end of the CNN lesson is a solid starting point. You'll have to change the number of classes and possibly the preprocessing, but aside from that it's plug and play!
#
# With the LeNet-5 solution from the lecture, you should expect a validation set accuracy of about 0.89. To meet specifications, the validation set accuracy will need to be at least 0.93. It is possible to get an even higher accuracy, but 0.93 is the minimum for a successful project submission.
#
# There are various aspects to consider when thinking about this problem:
#
# - Neural network architecture (is the network over or underfitting?)
# - Play around preprocessing techniques (normalization, rgb to grayscale, etc)
# - Number of examples per label (some have more than others).
# - Generate fake data.
#
# Here is an example of a [published baseline model on this problem](http://yann.lecun.com/exdb/publis/pdf/sermanet-ijcnn-11.pdf). It's not required to be familiar with the approach used in the paper but, it's good practice to try to read papers like these.

# ### Pre-process the Data Set (normalization, grayscale, etc.)

# Minimally, the image data should be normalized so that the data has mean zero and equal variance. For image data, `(pixel - 128)/ 128` is a quick way to approximately normalize the data and can be used in this project.
#
# Other pre-processing steps are optional. You can try different techniques to see if it improves performance.
#
# Use the code cell (or multiple code cells, if necessary) to implement the first step of your project.

# In[7]:

### Preprocess the data here. It is required to normalize the data. Other preprocessing steps could include
### converting to grayscale, etc.
import cv2

# Implement normalization of image data
def normalize(image_data):
    """
    Normalize the image data with
    :param image_data: The image data to be normalized
    :return: Normalized image data
    """
    if image_data[0].shape[2] == 3:
        grayscale = np.zeros(image_data.shape[:-1])
        for i in range(image_data.shape[0]):
            grayscale[i] = cv2.cvtColor(image_data[i], cv2.COLOR_RGB2GRAY)
        grayscale = np.expand_dims(grayscale, axis = 3)
        grayscale = (grayscale - 128) / 128
        image_data = grayscale
    return image_data

from sklearn.utils import shuffle

X_train, y_train = shuffle(X_train, y_train)

X_train_normalized = normalize(X_train)
X_valid_normalized = normalize(X_valid)
X_test_normalized = normalize(X_test)

#Update image shape
print(X_train[0].shape)
image_shape = X_train_normalized[0].shape
print(image_shape)

# ### Model Architecture

# In[8]:

### Define your architecture here.

import tensorflow as tf

EPOCHS = 25
BATCH_SIZE = 128

# ### LeNet layers

# In[9]:

from tensorflow.contrib.layers import flatten

def LeNet(x, input_depth = 1, n_classes = 43, keep_prob = tf.constant(1), layer_out = 5):
    # Arguments used for tf.truncated_normal, randomly defines variables for the weights and biases for each layer
    mu = 0
    sigma = 0.1

    # Weight and bias
    weights ={
        'layer_1': tf.Variable(tf.truncated_normal([5, 5, input_depth, 6], mu, sigma)),#6,16,120,84
        'layer_2': tf.Variable(tf.truncated_normal([5, 5, 6, 16], mu, sigma)),
        'layer_3': tf.Variable(tf.truncated_normal([400, 120], mu, sigma)),
        'layer_4': tf.Variable(tf.truncated_normal([120, 84], mu, sigma)),
        'layer_5': tf.Variable(tf.truncated_normal([84, n_classes], mu, sigma)),
    }
    biases = {
        'layer_1': tf.Variable(tf.zeros(6)),
        'layer_2': tf.Variable(tf.zeros(16)),
        'layer_3': tf.Variable(tf.zeros(120)),
        'layer_4': tf.Variable(tf.zeros(84)),
        'layer_5': tf.Variable(tf.zeros(n_classes)),
    }
    '''
    SOLUTION: Layer 1: Convolutional. Input = 32x32x3. Output = 28x28x18.
    SOLUTION: Activation.
    SOLUTION: Pooling. Input = 28x28x6. Output = 14x14x18.
    '''
    layer_1 = tf.nn.conv2d(x, weights['layer_1'], strides = [1, 1, 1, 1], padding = 'VALID')
    layer_1 = tf.nn.bias_add(layer_1, biases['layer_1'])
    layer_1 = tf.nn.relu(layer_1)
    layer_1 = tf.nn.max_pool(layer_1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
    #layer_1 = tf.nn.dropout(layer_1, keep_prob)

    '''
    SOLUTION: Layer 2: Convolutional. Output = 10x10x16.
    SOLUTION: Activation.
    SOLUTION: Pooling. Input = 10x10x16. Output = 5x5x16.
    '''
    layer_2 = tf.nn.conv2d(layer_1, weights['layer_2'], strides = [1, 1, 1, 1], padding = 'VALID')
    layer_2 = tf.nn.bias_add(layer_2, biases['layer_2'])
    layer_2 = tf.nn.relu(layer_2)
    layer_2 = tf.nn.max_pool(layer_2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')
    layer_2 = tf.nn.dropout(layer_2, keep_prob)

    '''
    SOLUTION: Flatten. Input = 5x5x16. Output = 400.
    '''
    flat_layer = flatten(layer_2)

    '''
    SOLUTION: Layer 3: Fully Connected. Input = 400. Output = 120.
    SOLUTION: Activation.
    '''
    layer_3 = tf.matmul(flat_layer, weights['layer_3']) + biases['layer_3']
    layer_3 = tf.nn.relu(layer_3)
    layer_3 = tf.nn.dropout(layer_3, keep_prob)

    '''
    SOLUTION: Layer 4: Fully Connected. Input = 120. Output = 84.
    SOLUTION: Activation.
    '''
    layer_4 = tf.matmul(layer_3, weights['layer_4']) + biases['layer_4']
    layer_4 = tf.nn.relu(layer_4)
    layer_4 = tf.nn.dropout(layer_4, keep_prob)

    '''
    SOLUTION: Layer 5: Fully Connected. Input = 84. Output = 10.
    '''
    layer_5 = tf.matmul(layer_4, weights['layer_5']) + biases['layer_5']
    logits = layer_5
    output = logits
    #decide which layer to output (just for the last optional point which plots the features)
    if layer_out == 1:
        output = layer_1
    elif layer_out == 2:
        output = layer_1
    elif layer_out == 3:
        output = layer_3
    elif layer_out == 4:
        output = layer_4
    return output

# ### Features and Label

# In[10]:

x = tf.placeholder(tf.float32, (None, image_shape[0], image_shape[1], image_shape[2]))
y = tf.placeholder(tf.int32, (None))
one_hot_y = tf.one_hot(y, n_classes)

# ### Train, Validate and Test the Model

# A validation set can be used to assess how well the model is performing. A low accuracy on the training and validation
# sets imply underfitting. A high accuracy on the training set but low accuracy on the validation set implies overfitting.

# In[11]:

### Train your model here.
### Calculate and report the accuracy on the training and validation set.
### Once a final model architecture is selected,
### the accuracy on the test set should be calculated and reported as well.
rate = 0.001
keep_prob = tf.placeholder(tf.float32)

logits = LeNet(x, keep_prob = keep_prob)
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(labels=one_hot_y, logits=logits)
loss_operation = tf.reduce_mean(cross_entropy)
optimizer = tf.train.AdamOptimizer(learning_rate = rate)
training_operation = optimizer.minimize(loss_operation)

### Model evaluation

correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(one_hot_y, 1))
accuracy_operation = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
saver = tf.train.Saver()

def evaluate(X_data, y_data):
    num_examples = len(X_data)
    total_accuracy = 0
    sess = tf.get_default_session()
    for offset in range(0, num_examples, BATCH_SIZE):
        batch_x, batch_y = X_data[offset:offset+BATCH_SIZE], y_data[offset:offset+BATCH_SIZE]
        accuracy = sess.run(accuracy_operation, feed_dict={x: batch_x, y: batch_y, keep_prob: 1})
        total_accuracy += (accuracy * len(batch_x))
    return total_accuracy / num_examples

# ## Train the Model
# Run the training data through the training pipeline to train the model.
#
# Before each epoch, shuffle the training set.
#
# After each epoch, measure the loss and accuracy of the validation set.
#
# Save the model after training.
#
# You do not need to modify this section.

# In[12]:

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    num_examples = len(X_train)

    print("Training...")
    print()
    for i in range(EPOCHS):
        X_train_normalized, y_train = shuffle(X_train_normalized, y_train)
        for offset in range(0, num_examples, BATCH_SIZE):
            end = offset + BATCH_SIZE
            batch_x, batch_y = X_train_normalized[offset:end], y_train[offset:end]
            sess.run(training_operation, feed_dict={x: batch_x, y: batch_y, keep_prob: 0.7})

        validation_accuracy = evaluate(X_valid_normalized, y_valid)
        print("EPOCH {} ...".format(i+1))
        print("Validation Accuracy = {:.3f}".format(validation_accuracy))
        print()

    saver.save(sess, './final_classifier')
    print("Model saved")

# ## Evaluate the Model
# Once you are completely satisfied with your model, evaluate the performance of the model on the test set.
#
# Be sure to only do this once!
#
# If you were to measure the performance of your trained model on the test set, then improve your model, and then measure the performance of your model on the test set again, that would invalidate your test results. You wouldn't get a true measure of how well your model would perform against real data.
#
# You do not need to modify this section.

# In[26]:

with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint('.'))

    test_accuracy = evaluate(X_test_normalized, y_test)
    print("Test Accuracy = {:.3f}".format(test_accuracy))

# ---
#
# ## Step 3: Test a Model on New Images
#
# To give yourself more insight into how your model is working, download at least five pictures of German traffic signs from the web and use your model to predict the traffic sign type.
#
# You may find `signnames.csv` useful as it contains mappings from the class id (integer) to the actual sign name.

# ### Load and Output the Images

# In[13]:

import matplotlib.image as mpimg

### Load the images and plot them here.
### Feel free to use as many code cells as needed.

y_personal = [9, 12, 14, 17, 1, 25]
X_personal = []

fig = plt.figure()
for i in range(len(y_personal)):
    axs = fig.add_subplot(1, 10, i + 1)
    axs.set_title(str([y_personal[i]]))
    file_name = '../images/newsign' + str(i + 1) + '.jpg'
    image = mpimg.imread(file_name)
    image = cv2.resize(image, (32, 32), interpolation = cv2.INTER_AREA)
    plt.imshow(image)
    X_personal.append(image)
X_personal = np.asarray(X_personal)
y_personal = np.asarray(y_personal)
plt.show()

print(X_personal.shape)
print(y_personal.shape)

# ### Predict the Sign Type for Each Image

# In[14]:

### Run the predictions here and use the model to output the prediction for each image.
### Make sure to pre-process the images with the same pre-processing pipeline used earlier.
### Feel free to use as many code cells as needed.
X_personal_normalized = normalize(X_personal)
prediction = tf.argmax(logits, 1)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, tf.train.latest_checkpoint('.'))
    sess = tf.get_default_session()
    print("Predicting...")
    print()
    #
    #input_x = np.expand_dims(X_personal_normalized, axis = 0)
    input_x = X_personal_normalized
    input_y = y_personal
    predicted = sess.run(prediction, feed_dict={x: input_x, y: input_y, keep_prob: 1})
    correct = sess.run(correct_prediction, feed_dict={x: input_x, y: input_y, keep_prob: 1})
    #print(predicted)
    #print(y_personal)
    #print(correct)
    for i in range(len(y_personal)):
        to_append = '[x]'
        if correct[i]:
            to_append = '[v]'
        print(to_append + ' Real = ' + str(sign_name_database[y_personal[i]]) + '   Predicted = ' + str(sign_name_database[predicted[i]]))

# ### Analyze Performance

# In[15]:

### Calculate the accuracy for these 5 new images.
### For example, if the model predicted 1 out of 5 signs correctly, it's 20% accurate on these new images.
with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint('.'))
    sess = tf.get_default_session()
    num_examples = len(X_personal_normalized)
    validation_accuracy = evaluate(X_personal_normalized, y_personal) * 100
    print("Validation Accuracy of personal images = " + str(int(validation_accuracy)) + '%')

# ### Output Top 5 Softmax Probabilities For Each Image Found on the Web

# For each of the new images, print out the model's softmax probabilities to show the **certainty** of the model's predictions (limit the output to the top 5 probabilities for each image). [`tf.nn.top_k`](https://www.tensorflow.org/versions/r0.12/api_docs/python/nn.html#top_k) could prove helpful here.
#
# The example below demonstrates how tf.nn.top_k can be used to find the top k predictions for each image.
#
# `tf.nn.top_k` will return the values and indices (class ids) of the top k predictions. So if k=3, for each sign, it'll return the 3 largest probabilities (out of a possible 43) and the correspoding class ids.
#
# Take this numpy array as an example. The values in the array represent predictions. The array contains softmax probabilities for five candidate images with six possible classes. `tf.nn.top_k` is used to choose the three classes with the highest probability:
#
# ```
# # (5, 6) array
# a = np.array([[ 0.24879643,  0.07032244,  0.12641572,  0.34763842,  0.07893497,
#          0.12789202],
#        [ 0.28086119,  0.27569815,  0.08594638,  0.0178669 ,  0.18063401,
#          0.15899337],
#        [ 0.26076848,  0.23664738,  0.08020603,  0.07001922,  0.1134371 ,
#          0.23892179],
#        [ 0.11943333,  0.29198961,  0.02605103,  0.26234032,  0.1351348 ,
#          0.16505091],
#        [ 0.09561176,  0.34396535,  0.0643941 ,  0.16240774,  0.24206137,
#          0.09155967]])
# ```
#
# Running it through `sess.run(tf.nn.top_k(tf.constant(a), k=3))` produces:
#
# ```
# TopKV2(values=array([[ 0.34763842,  0.24879643,  0.12789202],
#        [ 0.28086119,  0.27569815,  0.18063401],
#        [ 0.26076848,  0.23892179,  0.23664738],
#        [ 0.29198961,  0.26234032,  0.16505091],
#        [ 0.34396535,  0.24206137,  0.16240774]]), indices=array([[3, 0, 5],
#        [0, 1, 4],
#        [0, 5, 1],
#        [1, 3, 5],
#        [1, 4, 3]], dtype=int32))
# ```
#
# Looking just at the first row we get `[ 0.34763842,  0.24879643,  0.12789202]`, you can confirm these are the 3 largest probabilities in `a`. You'll also notice `[3, 0, 5]` are the corresponding indices.

# In[16]:

### Print out the top five softmax probabilities for the predictions on the German traffic sign images found on the web.
### Feel free to use as many code cells as needed.
top5 = tf.nn.top_k(tf.nn.softmax(logits), k = 5, sorted = True)
with tf.Session() as sess:
    saver.restore(sess, tf.train.latest_checkpoint('.'))
    sess = tf.get_default_session()

    for i in range(len(y_personal)):
        input_x = np.expand_dims(X_personal_normalized[i], axis = 0)
        input_y = y_personal
        output_top5 = sess.run(top5, feed_dict={x: input_x, y: input_y, keep_prob: 1})
        print()
        print('Input image = ' + sign_name_database[y_personal[i]]   + '. Possiblities are:')
        for j in range(len(output_top5[1][0])):
            print(sign_name_database[output_top5[1][0][j]] + ' with probability ' + str(output_top5[0][0][j]))

# ### Project Writeup
#
# Once you have completed the code implementation, document your results in a project writeup using this [template](https://github.com/udacity/CarND-Traffic-Sign-Classifier-Project/blob/master/writeup_template.md) as a guide. The writeup can be in a markdown or pdf file.

# > **Note**: Once you have completed all of the code implementations and successfully answered each question above, you may finalize your work by exporting the iPython Notebook as an HTML document. You can do this by using the menu above and navigating to  \n",
#     "**File -> Download as -> HTML (.html)**. Include the finished document along with this notebook as your submission.

# ---
#
# ## Step 4 (Optional): Visualize the Neural Network's State with Test Images
#
#  This Section is not required to complete but acts as an additional excersise for understaning the output of a neural network's weights. While neural networks can be a great learning device they are often referred to as a black box. We can understand what the weights of a neural network look like better by plotting their feature maps. After successfully training your neural network you can see what it's feature maps look like by plotting the output of the network's weight layers in response to a test stimuli image. From these plotted feature maps, it's possible to see what characteristics of an image the network finds interesting. For a sign, maybe the inner network feature maps react with high activation to the sign's boundary outline or to the contrast in the sign's painted symbol.
#
#  Provided for you below is the function code that allows you to get the visualization output of any tensorflow weight layer you want. The inputs to the function should be a stimuli image, one used during training or a new one you provided, and then the tensorflow variable name that represents the layer's state during the training process, for instance if you wanted to see what the [LeNet lab's](https://classroom.udacity.com/nanodegrees/nd013/parts/fbf77062-5703-404e-b60c-95b78b2f3f9e/modules/6df7ae49-c61c-4bb2-a23e-6527e69209ec/lessons/601ae704-1035-4287-8b11-e2c2716217ad/concepts/d4aca031-508f-4e0b-b493-e7b706120f81) feature maps looked like for it's second convolutional layer you could enter conv2 as the tf_activation variable.
#
# For an example of what feature map outputs look like, check out NVIDIA's results in their paper [End-to-End Deep Learning for Self-Driving Cars](https://devblogs.nvidia.com/parallelforall/deep-learning-self-driving-cars/) in the section Visualization of internal CNN State. NVIDIA was able to show that their network's inner weights had high activations to road boundary lines by comparing feature maps from an image with a clear path to one without. Try experimenting with a similar test to show that your trained network's weights are looking for interesting features, whether it's looking at differences in feature maps from images with or without a sign, or even what feature maps look like in a trained network vs a completely untrained one on the same sign image.
#
# <figure>
#  <img src="../images/visualize_cnn.png" width="380" alt="Combined Image" />
#  <figcaption>
#  <p></p>
#  <p style="text-align: center;"> Your output should look something like this (above)</p>
#  </figcaption>
# </figure>
#  <p></p>
#

# In[17]:

### Visualize your network's feature maps here.
### Feel free to use as many code cells as needed.

# image_input: the test image being fed into the network to produce the feature maps
# tf_activation: should be a tf variable name used during your training procedure that represents the calculated state of a specific weight layer
# activation_min/max: can be used to view the activation contrast in more detail, by default matplot sets min and max to the actual min and max values of the output
# plt_num: used to plot out multiple different weight feature map sets on the same block, just extend the plt number for each new feature map entry

def outputFeatureMap(image_input, tf_activation, activation_min=-1, activation_max=-1 ,plt_num=1):
    # Here make sure to preprocess your image_input in a way your network expects
    # with size, normalization, ect if needed
    # image_input =
    # Note: x should be the same name as your network's tensorflow data placeholder variable
    # If you get an error tf_activation is not defined it may be having trouble accessing the variable from inside a function
    activation = tf_activation.eval(session=sess,feed_dict={x : image_input})
    featuremaps = activation.shape[3]
    plt.figure(plt_num, figsize=(10,10))
    for featuremap in range(featuremaps):
        plt.subplot(6,8, featuremap+1) # sets the number of feature maps to show on each row and column
        plt.title('FeatureMap ' + str(featuremap)) # displays the feature map number
        if activation_min != -1 & activation_max != -1:
            plt.imshow(activation[0,:,:, featuremap], interpolation="nearest", vmin =activation_min, vmax=activation_max, cmap="gray")
        elif activation_max != -1:
            plt.imshow(activation[0,:,:, featuremap], interpolation="nearest", vmax=activation_max, cmap="gray")
        elif activation_min !=-1:
            plt.imshow(activation[0,:,:, featuremap], interpolation="nearest", vmin=activation_min, cmap="gray")
        else:
            plt.imshow(activation[0,:,:, featuremap], interpolation="nearest", cmap="gray")

# In[25]:

index = random.randint(0, len(X_train))
image = ((X_train_normalized[index] * 128) + 128).squeeze()

plt.figure(figsize=(1,1))
plt.title(str(y_train[index]) +' = ' +  sign_name_database[y_train[index]])
plt.imshow(image)
plt.show()

tf_activation_1 = LeNet(x, layer_out = 1, keep_prob = keep_prob)
tf_activation_2 = LeNet(x, layer_out = 2, keep_prob = keep_prob)
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, tf.train.latest_checkpoint('.'))
    sess = tf.get_default_session()

    input_x = np.expand_dims(X_train_normalized[index], axis = 0)
    input_y = y_train[index]

    outputFeatureMap(input_x, tf_activation_1, plt_num = 1)
    outputFeatureMap(input_x, tf_activation_2, plt_num = 2)

# In[ ]:

# In[ ]:


import numpy as np
import warnings
import sys
try:
    import pycuda.autoinit
    import pycuda.driver as drv
    from pycuda import gpuarray
    from pycuda.compiler import SourceModule
    import pycuda.cumath
    pycuda_available = True
except:
    pycuda_available = False

use_gpu = False

default_blocksize = 128

def enable_gpu(enable=True):
    """Sets the use_gpu flag to enable/disable the use of CUDA kernels.

    Args:
        enable (bool): Set use_gpu flag to this value (default=True).
    """
    # todo: do gpuarrays require cuda toolkit? otherwise distinguish if only gpuarrays work but no kernel compilation
    global use_gpu

    if use_gpu == enable:
        return

    if not enable:
        use_gpu = False
        sys.stdout.write('Disabling GPU usage.\n')
        sys.stdout.flush()
        return

    if not pycuda_available:
        warnings.warn("Unable to import PyCuda - fall back to CPU mode")
        use_gpu = False
        sys.stdout.write('Disabling GPU usage.\n')
        sys.stdout.flush()
    else:
        try:
            test_fun = SourceModule("""
                __global__ void test_kernel(float *x)
                {
                    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
                    x[i] = 2 * i;
                }""").get_function("test_kernel")
            x = gpuarray.to_gpu(np.zeros(128, dtype=np.float32))
            test_fun(x, block=(128,1,1), grid=(1,1))
            assert np.allclose(x.get(), 2 * np.arange(128))
            use_gpu = True
            sys.stdout.write('Enabling GPU usage.\n')
            sys.stdout.flush()
            return

        except Exception as e:
            warnings.warn("CUDA test kernel failed - fall back to CPU mode")
            print(e)
            use_gpu = False
            sys.stdout.write('Disabling GPU usage.\n')
            sys.stdout.flush()

# The following cuda kernel multiplies the coupling matrix to a vector. It is based on linear interpolation of
# the lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# z_pos (np.float32): z_pos[i] contains the respective particle z-position
# re_lookup_pl (np.float32): the real part of the lookup table for the z1+z2 part of the Sommerfeld integral,
#                            in the format (rho, sum_z, n1, n2)
# im_lookup_pl (np.float32): the imaginary part of the lookup table for the z1+z2 part of the Sommerfeld
#                            integral, in the format (rho, sum_z, n1, n2)
# re_lookup_mn (np.float32): the real part of the lookup table for the z1-z2 part of the Sommerfeld integral,
#                            in the format (rho, diff_z, n1, n2)
# im_lookup_mn (np.float32): the imaginary part of the lookup table for the z1-z2 part of the Sommerfeld
#                            integral, in the format (rho, diff_z, n1, n2)
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
linear_volume_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define Z_ARRAY_LENGTH %i
    #define MIN_RHO %f
    #define MIN_Z_SUM %f
    #define MIN_Z_DIFF %f
    #define LOOKUP_RESOLUTION %f
    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *z_pos, const float *re_lookup_pl, const float *im_lookup_pl,
                                    const float *re_lookup_mn, const float *im_lookup_mn,
                                    const float *re_in_vec, const float *im_in_vec,
                                    float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];
        const float z1 = z_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];
            float sz21 = z1 + z_pos[i2];
            float dz21 = z1 - z_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float rho = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int rho_idx = (int) floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);
            float rho_w = (rho - MIN_RHO) / LOOKUP_RESOLUTION - floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);

            int sz_idx = (int) floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);
            float sz_w = (sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION - floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);

            int dz_idx = (int) floor((dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION);
            float dz_w = (dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION
                            - floor((dz21 - MIN_Z_DIFF) /  LOOKUP_RESOLUTION);

            int idx_rho_sz = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                +  sz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rho_szpl1 = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + (sz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_sz = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + sz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_szpl1 = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                    + (sz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rho_dz = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE + dz_idx * BLOCKSIZE * BLOCKSIZE
                                + n1 * BLOCKSIZE + n2;

            int idx_rho_dzpl1 = rho_idx * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + (dz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_dz = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                + dz_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            int idx_rhopl1_dzpl1 = (rho_idx + 1) * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE
                                    + (dz_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float f_rhoi = re_lookup_pl[idx_rho_sz] * (1 - sz_w) + re_lookup_pl[idx_rho_szpl1] * sz_w;
            float f_rhoipl1 = re_lookup_pl[idx_rhopl1_sz] * (1 - sz_w) + re_lookup_pl[idx_rhopl1_szpl1] * sz_w;
            float re_si_pl = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = im_lookup_pl[idx_rho_sz] * (1 - sz_w) + im_lookup_pl[idx_rho_szpl1] * sz_w;
            f_rhoipl1 = im_lookup_pl[idx_rhopl1_sz] * (1 - sz_w) + im_lookup_pl[idx_rhopl1_szpl1] * sz_w;
            float im_si_pl = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = re_lookup_mn[idx_rho_dz] * (1 - dz_w) + re_lookup_mn[idx_rho_dzpl1] * dz_w;
            f_rhoipl1 = re_lookup_mn[idx_rhopl1_dz] * (1 - dz_w) + re_lookup_mn[idx_rhopl1_dzpl1] * dz_w;
            float re_si_mn = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            f_rhoi = im_lookup_mn[idx_rho_dz] * (1 - dz_w) + im_lookup_mn[idx_rho_dzpl1] * dz_w;
            f_rhoipl1 = im_lookup_mn[idx_rhopl1_dz] * (1 - dz_w) + im_lookup_mn[idx_rhopl1_dzpl1] * dz_w;
            float im_si_mn = f_rhoi * (1 - rho_w) + f_rhoipl1 * rho_w;

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * (re_si_pl + re_si_mn) - im_eimphi * (im_si_pl + im_si_mn);
            float im_w = im_eimphi * (re_si_pl + re_si_mn) + re_eimphi * (im_si_pl + im_si_mn);

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# The following cuda kernel multiplies the coupling matrix to a vector. It is based on cubic Hermite spline interpolation
# of the lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# z_pos (np.float32): z_pos[i] contains the respective particle z-position
# re_lookup_pl (np.float32): the real part of the lookup table for the z1+z2 part of the Sommerfeld integral,
#                            in the format (rho, sum_z, n1, n2)
# im_lookup_pl (np.float32): the imaginary part of the lookup table for the z1+z2 part of the Sommerfeld
#                            integral, in the format (rho, sum_z, n1, n2)
# re_lookup_mn (np.float32): the real part of the lookup table for the z1-z2 part of the Sommerfeld integral,
#                            in the format (rho, diff_z, n1, n2)
# im_lookup_mn (np.float32): the imaginary part of the lookup table for the z1-z2 part of the Sommerfeld
#                            integral, in the format (rho, diff_z, n1, n2)
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
cubic_volume_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define Z_ARRAY_LENGTH %i
    #define MIN_RHO %f
    #define MIN_Z_SUM %f
    #define MIN_Z_DIFF %f
    #define LOOKUP_RESOLUTION %f

    __device__ int lookup_index(int const i_r, int const i_z, int const n1, int const n2)
    {
        return i_r * Z_ARRAY_LENGTH * BLOCKSIZE * BLOCKSIZE + i_z * BLOCKSIZE * BLOCKSIZE + n1 *  BLOCKSIZE + n2;
    }

    __device__ float cubic_interpolation(float const w, float const lookup_imn1, float const lookup_i,
                                            float const lookup_ipl1, float const lookup_ipl2)
    {
        return ((-w*w*w+2*w*w-w) * lookup_imn1 + (3*w*w*w-5*w*w+2) * lookup_i + (-3*w*w*w+4*w*w+w) * lookup_ipl1
                + (w*w*w-w*w) * lookup_ipl2) / 2;
    }

    __device__ float lookup_2D(int const i_rho, float const w_rho, int const i_z, float const w_z, int const n1,
                                int const n2, float const *lookup)
    {
        int i_zmn1 = lookup_index(i_rho-1, i_z-1, n1, n2);
        int i = lookup_index(i_rho-1, i_z, n1, n2);
        int i_zpl1 = lookup_index(i_rho-1, i_z+1, n1, n2);
        int i_zpl2 = lookup_index(i_rho-1, i_z+2, n1, n2);
        float lookup_rmn1 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho, i_z-1, n1, n2);
        i = lookup_index(i_rho, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho, i_z+2, n1, n2);
        float lookup_r = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho+1, i_z-1, n1, n2);
        i = lookup_index(i_rho+1, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho+1, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho+1, i_z+2, n1, n2);
        float lookup_rpl1 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        i_zmn1 = lookup_index(i_rho+2, i_z-1, n1, n2);
        i = lookup_index(i_rho+2, i_z, n1, n2);
        i_zpl1 = lookup_index(i_rho+2, i_z+1, n1, n2);
        i_zpl2 = lookup_index(i_rho+2, i_z+2, n1, n2);
        float lookup_rpl2 = cubic_interpolation(w_z, lookup[i_zmn1], lookup[i], lookup[i_zpl1], lookup[i_zpl2]);

        return cubic_interpolation(w_rho, lookup_rmn1, lookup_r, lookup_rpl1, lookup_rpl2);
    }

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *z_pos, const float *re_lookup_pl, const float *im_lookup_pl,
                                    const float *re_lookup_mn, const float *im_lookup_mn,
                                    const float *re_in_vec, const float *im_in_vec,
                                    float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];
        const float z1 = z_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];
            float sz21 = z1 + z_pos[i2];
            float dz21 = z1 - z_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float rho = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int rho_idx = (int) floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);
            float rho_w = (rho - MIN_RHO) / LOOKUP_RESOLUTION - floor((rho - MIN_RHO) / LOOKUP_RESOLUTION);

            int sz_idx = (int) floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);
            float sz_w = (sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION - floor((sz21 - MIN_Z_SUM) / LOOKUP_RESOLUTION);

            int dz_idx = (int) floor((dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION);
            float dz_w = (dz21 - MIN_Z_DIFF) / LOOKUP_RESOLUTION
                            - floor((dz21 - MIN_Z_DIFF) /  LOOKUP_RESOLUTION);

            float re_si_pl = lookup_2D(rho_idx, rho_w, sz_idx, sz_w, n1, n2, re_lookup_pl);
            float im_si_pl = lookup_2D(rho_idx, rho_w, sz_idx, sz_w, n1, n2, im_lookup_pl);
            float re_si_mn = lookup_2D(rho_idx, rho_w, dz_idx, dz_w, n1, n2, re_lookup_mn);
            float im_si_mn = lookup_2D(rho_idx, rho_w, dz_idx, dz_w, n1, n2, im_lookup_mn);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * (re_si_pl + re_si_mn) - im_eimphi * (im_si_pl + im_si_mn);
            float im_w = im_eimphi * (re_si_pl + re_si_mn) + re_eimphi * (im_si_pl + im_si_mn);

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# This cuda kernel multiplies the coupling matrix to a vector. It is based on linear interpolation of the lookup
# table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# re_lookup (np.float32): the real part of the lookup table, in the format [r, n1, n2]
# im_lookup (np.float32): the imaginary part of the lookup table, in the format [r, n1, n2]
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
linear_radial_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define MIN_RHO %f
    #define LOOKUP_RESOLUTION %f

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *re_lookup, const float *im_lookup, const float *re_in_vec,
                                    const float *im_in_vec, float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float r = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int r_idx = (int) floor((r - MIN_RHO) / LOOKUP_RESOLUTION);
            float w = (r - MIN_RHO) / LOOKUP_RESOLUTION - floor((r - MIN_RHO) / LOOKUP_RESOLUTION);

            int idx = r_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl_1 = (r_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float re_si = re_lookup[idx] + w * (re_lookup[idx_pl_1] - re_lookup[idx]);
            float im_si = im_lookup[idx] + w * (im_lookup[idx_pl_1] - im_lookup[idx]);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * re_si - im_eimphi * im_si;
            float im_w = im_eimphi * re_si + re_eimphi * im_si;

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];

        }
    }"""

# This cuda kernel multiplies the coupling matrix to a vector. It is based on cubic Hermite spline interpolation of the
# lookup table.
#
# input arguments of the cuda kernel:
# n (np.uint32):  n[i] contains the mutlipole multi-index with regard to self.l_max and self.m_max of the i-th
#                 entry of a system vector
# m (np.float32): m[i] contains the multipole order with regard to particle.l_max and particle.m_max
# x_pos (np.float32): x_pos[i] contains the respective particle x-position
# y_pos (np.float32): y_pos[i] contains the respective particle y-position
# re_lookup (np.float32): the real part of the lookup table, in the format [r, n1, n2]
# im_lookup (np.float32): the imaginary part of the lookup table, in the format [r, n1, n2]
# re_in_vec (np.float32): the real part of the vector to be multiplied with the coupling matrix
# im_in_vec (np.float32): the imaginary part of the vector to be multiplied with the coupling matrix
# re_result_vec (np.float32): the real part of the vector into which the result is written
# im_result_vec (np.float32): the imaginary part of the vector into which the result is written
cubic_radial_lookup_source = """
    #define BLOCKSIZE %i
    #define NUMBER_OF_UNKNOWNS %i
    #define MIN_RHO %f
    #define LOOKUP_RESOLUTION %f

    __device__ float cubic_interpolation(float const w, float const lookup_imn1, float const lookup_i,
                                         float const lookup_ipl1, float const lookup_ipl2)
    {
        return ((-w*w*w+2*w*w-w) * lookup_imn1 + (3*w*w*w-5*w*w+2) * lookup_i + (-3*w*w*w+4*w*w+w) * lookup_ipl1
                + (w*w*w-w*w) * lookup_ipl2) / 2;
    }

    __global__ void coupling_kernel(const int *n, const float *m, const float *x_pos, const float *y_pos,
                                    const float *re_lookup, const float *im_lookup, const float *re_in_vec,
                                    const float *im_in_vec, float  *re_result, float  *im_result)
    {
        unsigned int i1 = blockIdx.x * blockDim.x + threadIdx.x;
        if(i1 >= NUMBER_OF_UNKNOWNS) return;

        const float x1 = x_pos[i1];
        const float y1 = y_pos[i1];

        const int n1 = n[i1];
        const float m1 = m[i1];

        re_result[i1] = 0.0;
        im_result[i1] = 0.0;

        for (int i2=0; i2<NUMBER_OF_UNKNOWNS; i2++)
        {
            float x21 = x1 - x_pos[i2];
            float y21 = y1 - y_pos[i2];

            const int n2 = n[i2];
            const float m2 = m[i2];

            float r = sqrt(x21*x21+y21*y21);
            float phi = atan2(y21,x21);

            int r_idx = (int) floor((r - MIN_RHO) / LOOKUP_RESOLUTION);
            float w = (r - MIN_RHO) / LOOKUP_RESOLUTION - floor((r - MIN_RHO) / LOOKUP_RESOLUTION);

            int idx_mn1 = (r_idx - 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx = r_idx * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl1 = (r_idx + 1) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;
            int idx_pl2 = (r_idx + 2) * BLOCKSIZE * BLOCKSIZE + n1 * BLOCKSIZE + n2;

            float re_si = cubic_interpolation(w, re_lookup[idx_mn1], re_lookup[idx], re_lookup[idx_pl1],
                                              re_lookup[idx_pl2]);
            float im_si = cubic_interpolation(w, im_lookup[idx_mn1], im_lookup[idx], im_lookup[idx_pl1],
                                              im_lookup[idx_pl2]);

            float re_eimphi = cosf((m2 - m1) * phi);
            float im_eimphi = sinf((m2 - m1) * phi);

            float re_w = re_eimphi * re_si - im_eimphi * im_si;
            float im_w = im_eimphi * re_si + re_eimphi * im_si;

            re_result[i1] += re_w * re_in_vec[i2] - im_w * im_in_vec[i2];
            im_result[i1] += re_w * im_in_vec[i2] + im_w * re_in_vec[i2];
        }
    }"""

# This cuda kernel is used for the calculation of volume lookup tables (see smuthi.particle_coupling module).
volume_lookup_assembly_code = """
    #define BLOCKSIZE %i
    #define RHO_ARRAY_LENGTH %i
    #define Z_ARRAY_LENGTH %i
    #define K_ARRAY_LENGTH %i

    __global__ void helper(const float *re_bes_jac, const float *im_bes_jac, const float *re_belbee,
                            const float *im_belbee, const float *re_d_kappa, const float *im_d_kappa,
                            float *re_result, float  *im_result)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if(i >= RHO_ARRAY_LENGTH * Z_ARRAY_LENGTH) return;

        unsigned int i_rho = i / Z_ARRAY_LENGTH;
        unsigned int i_z = i %% Z_ARRAY_LENGTH;

        float re_res = 0.0;
        float im_res = 0.0;

        int i_kr = i_rho * K_ARRAY_LENGTH;
        int i_kz = i_z * K_ARRAY_LENGTH;

        float re_integrand_kp1 = re_bes_jac[i_kr] * re_belbee[i_kz] - im_bes_jac[i_kr] * im_belbee[i_kz];
        float im_integrand_kp1 = re_bes_jac[i_kr] * im_belbee[i_kz] + im_bes_jac[i_kr] * re_belbee[i_kz];

        for (int i_k=0; i_k<(K_ARRAY_LENGTH-1); i_k++)
        {
            i_kr = i_rho * K_ARRAY_LENGTH + i_k;
            i_kz = i_z * K_ARRAY_LENGTH + i_k;

            float re_integrand = re_integrand_kp1;
            float im_integrand = im_integrand_kp1;

            re_integrand_kp1 = re_bes_jac[i_kr+1] * re_belbee[i_kz+1] - im_bes_jac[i_kr+1] * im_belbee[i_kz+1];
            im_integrand_kp1 = re_bes_jac[i_kr+1] * im_belbee[i_kz+1] + im_bes_jac[i_kr+1] * re_belbee[i_kz+1];

            float re_sint = re_integrand + re_integrand_kp1;
            float im_sint = im_integrand + im_integrand_kp1;

            re_res += 0.5 * (re_sint * re_d_kappa[i_k] - im_sint * im_d_kappa[i_k]);
            im_res += 0.5 * (re_sint * im_d_kappa[i_k] + im_sint * re_d_kappa[i_k]);
        }

        re_result[i] = re_res;
        im_result[i] = im_res;

    }"""

# This cuda kernel is used for the calculation of radial lookup tables (see smuthi.particle_coupling module).
radial_lookup_assembly_code = """
    #define BLOCKSIZE %i
    #define RHO_ARRAY_LENGTH %i
    #define K_ARRAY_LENGTH %i

    __global__ void helper(const float *re_bes_jac, const float *im_bes_jac, const float *re_belbee,
                            const float *im_belbee, const float *re_d_kappa, const float *im_d_kappa,
                            float *re_result, float  *im_result)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if(i >= RHO_ARRAY_LENGTH) return;

        float re_res = 0.0;
        float im_res = 0.0;

        int i_kr = i * K_ARRAY_LENGTH;

        float re_integrand_kp1 = re_bes_jac[i_kr] * re_belbee[0] - im_bes_jac[i_kr] * im_belbee[0];
        float im_integrand_kp1 = re_bes_jac[i_kr] * im_belbee[0] + im_bes_jac[i_kr] * re_belbee[0];

        for (int i_k=0; i_k<(K_ARRAY_LENGTH-1); i_k++)
        {
            i_kr = i * K_ARRAY_LENGTH + i_k;

            float re_integrand = re_integrand_kp1;
            float im_integrand = im_integrand_kp1;

            re_integrand_kp1 = re_bes_jac[i_kr+1] * re_belbee[i_k+1] - im_bes_jac[i_kr+1] * im_belbee[i_k+1];
            im_integrand_kp1 = re_bes_jac[i_kr+1] * im_belbee[i_k+1] + im_bes_jac[i_kr+1] * re_belbee[i_k+1];

            float re_sint = re_integrand + re_integrand_kp1;
            float im_sint = im_integrand + im_integrand_kp1;

            re_res += 0.5 * (re_sint * re_d_kappa[i_k] - im_sint * im_d_kappa[i_k]);
            im_res += 0.5 * (re_sint * im_d_kappa[i_k] + im_sint * re_d_kappa[i_k]);
        }

        re_result[i] = re_res;
        im_result[i] = im_res;

    }"""

# This cuda kernel is used for the evaluation of the electric field of plane wave expansions.
pwe_electric_field_evaluation_code = """
    #define LEN_X %i
    #define LEN_K %i
    #define LEN_A %i
    #define RE_ONE_OVER_K %.10f
    #define IM_ONE_OVER_K %.10f

    __device__ void alpha_integral(const int i_k, const float x, const float y, const float z,
                                   const float re_kp, const float im_kp, const float re_kz, const float im_kz,
                                   const float *alpha_array,
                                   const float *re_g_te_array, const float *im_g_te_array,
                                   const float *re_g_tm_array, const float *im_g_tm_array,
                                   float *re_integral_x, float *im_integral_x,
                                   float *re_integral_y, float *im_integral_y,
                                   float *re_integral_z, float *im_integral_z)
    {
        float a = 0.0;

        float re_x_result = 0.0;
        float im_x_result = 0.0;
        float re_y_result = 0.0;
        float im_y_result = 0.0;
        float re_z_result = 0.0;
        float im_z_result = 0.0;

        float re_e_x_alpha_integrand = 0.0;
        float im_e_x_alpha_integrand = 0.0;
        float re_e_y_alpha_integrand = 0.0;
        float im_e_y_alpha_integrand = 0.0;
        float re_e_z_alpha_integrand = 0.0;
        float im_e_z_alpha_integrand = 0.0;

        for (int i_a=0; i_a<LEN_A; i_a++)
        {
            float a_old = a;
            a = alpha_array[i_a];
            float d_a = a - a_old;
            float sin_a = sinf(a);
            float cos_a = cosf(a);

            float re_ikr = -im_kp * (cos_a * x + sin_a * y) - im_kz * z;
            float im_ikr = re_kp * (cos_a * x + sin_a * y) + re_kz * z;

            float ereikr = expf(re_ikr);
            float re_eikr =  ereikr * cosf(im_ikr);
            float im_eikr =  ereikr * sinf(im_ikr);

            // integrand
            float re_e_x_alpha_integrand_old = re_e_x_alpha_integrand;
            float im_e_x_alpha_integrand_old = im_e_x_alpha_integrand;
            float re_e_y_alpha_integrand_old = re_e_y_alpha_integrand;
            float im_e_y_alpha_integrand_old = im_e_y_alpha_integrand;
            float re_e_z_alpha_integrand_old = re_e_z_alpha_integrand;
            float im_e_z_alpha_integrand_old = im_e_z_alpha_integrand;

            int i_ka = i_k * LEN_A + i_a;

            // pol=TE
            float re_g = re_g_te_array[i_ka];
            float im_g = im_g_te_array[i_ka];

            float re_geikr = re_g * re_eikr - im_g * im_eikr;
            float im_geikr = re_g * im_eikr + im_g * re_eikr;

            re_e_x_alpha_integrand = -sin_a * re_geikr;
            im_e_x_alpha_integrand = -sin_a * im_geikr;
            re_e_y_alpha_integrand = cos_a * re_geikr;
            im_e_y_alpha_integrand = cos_a * im_geikr;

            // pol=TM
            re_g = re_g_tm_array[i_ka];
            im_g = im_g_tm_array[i_ka];

            re_geikr = re_g * re_eikr - im_g * im_eikr;
            im_geikr = re_g * im_eikr + im_g * re_eikr;

            float re_kzgeikr = re_kz * re_geikr - im_kz * im_geikr;
            float im_kzgeikr = re_kz * im_geikr + im_kz * re_geikr;

            float re_kzkgeikr = re_kzgeikr * RE_ONE_OVER_K - im_kzgeikr * IM_ONE_OVER_K;
            float im_kzkgeikr = re_kzgeikr * IM_ONE_OVER_K + im_kzgeikr * RE_ONE_OVER_K;

            re_e_x_alpha_integrand += cos_a * re_kzkgeikr;
            im_e_x_alpha_integrand += cos_a * im_kzkgeikr;
            re_e_y_alpha_integrand += sin_a * re_kzkgeikr;
            im_e_y_alpha_integrand += sin_a * im_kzkgeikr;

            float re_kpgeikr = re_kp * re_geikr - im_kp * im_geikr;
            float im_kpgeikr = re_kp * im_geikr + im_kp * re_geikr;
            re_e_z_alpha_integrand = -(re_kpgeikr * RE_ONE_OVER_K - im_kpgeikr * IM_ONE_OVER_K);
            im_e_z_alpha_integrand = -(re_kpgeikr * IM_ONE_OVER_K + im_kpgeikr * RE_ONE_OVER_K);

            if (i_a>0)
            {
                re_x_result += 0.5 * d_a * (re_e_x_alpha_integrand + re_e_x_alpha_integrand_old);
                im_x_result += 0.5 * d_a * (im_e_x_alpha_integrand + im_e_x_alpha_integrand_old);
                re_y_result += 0.5 * d_a * (re_e_y_alpha_integrand + re_e_y_alpha_integrand_old);
                im_y_result += 0.5 * d_a * (im_e_y_alpha_integrand + im_e_y_alpha_integrand_old);
                re_z_result += 0.5 * d_a * (re_e_z_alpha_integrand + re_e_z_alpha_integrand_old);
                im_z_result += 0.5 * d_a * (im_e_z_alpha_integrand + im_e_z_alpha_integrand_old);
            }
        }
        re_integral_x[0] = re_x_result;
        im_integral_x[0] = im_x_result;
        re_integral_y[0] = re_y_result;
        im_integral_y[0] = im_y_result;
        re_integral_z[0] = re_z_result;
        im_integral_z[0] = im_z_result;
    }

    __global__ void electric_field(const float *re_kp_array, const float *im_kp_array,
                                   const float *re_kz_array, const float *im_kz_array,
                                   const float *alpha_array,
                                   const float *x_array, const float *y_array, const float *z_array,
                                   const float *re_g_te_array, const float *im_g_te_array,
                                   const float *re_g_tm_array, const float *im_g_tm_array,
                                   float *re_e_x, float *im_e_x, float *re_e_y, float *im_e_y,
                                   float *re_e_z, float *im_e_z)
    {
        unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= LEN_X) return;

        float x = x_array[i];
        float y = y_array[i];
        float z = z_array[i];

        float re_kp = 0.0;
        float im_kp = 0.0;

        float re_x_result = 0.0;
        float im_x_result = 0.0;
        float re_y_result = 0.0;
        float im_y_result = 0.0;
        float re_z_result = 0.0;
        float im_z_result = 0.0;

        float re_x_kappa_integrand = 0.0;
        float im_x_kappa_integrand = 0.0;
        float re_y_kappa_integrand = 0.0;
        float im_y_kappa_integrand = 0.0;
        float re_z_kappa_integrand = 0.0;
        float im_z_kappa_integrand = 0.0;

        for (int i_k=0; i_k<LEN_K; i_k++)
        {
            float re_kp_old = re_kp;
            float im_kp_old = im_kp;

            re_kp = re_kp_array[i_k];
            im_kp = im_kp_array[i_k];

            float re_d_kp = re_kp - re_kp_old;
            float im_d_kp = im_kp - im_kp_old;

            float re_kz = re_kz_array[i_k];
            float im_kz = im_kz_array[i_k];

            float re_x_kappa_integrand_old = re_x_kappa_integrand;
            float im_x_kappa_integrand_old = im_x_kappa_integrand;
            float re_y_kappa_integrand_old = re_y_kappa_integrand;
            float im_y_kappa_integrand_old = im_y_kappa_integrand;
            float re_z_kappa_integrand_old = re_z_kappa_integrand;
            float im_z_kappa_integrand_old = im_z_kappa_integrand;

            float re_x_alpha_itgrl = 0.0;
            float im_x_alpha_itgrl = 0.0;
            float re_y_alpha_itgrl = 0.0;
            float im_y_alpha_itgrl = 0.0;
            float re_z_alpha_itgrl = 0.0;
            float im_z_alpha_itgrl = 0.0;

            alpha_integral(i_k, x, y, z, re_kp, im_kp, re_kz, im_kz, alpha_array,
                           re_g_te_array, im_g_te_array, re_g_tm_array, im_g_tm_array,
                           &re_x_alpha_itgrl, &im_x_alpha_itgrl, &re_y_alpha_itgrl, &im_y_alpha_itgrl,
                           &re_z_alpha_itgrl, &im_z_alpha_itgrl);

            re_x_kappa_integrand = re_x_alpha_itgrl * re_kp - im_x_alpha_itgrl * im_kp;
            im_x_kappa_integrand = re_x_alpha_itgrl * im_kp + im_x_alpha_itgrl * re_kp;
            re_y_kappa_integrand = re_y_alpha_itgrl * re_kp - im_y_alpha_itgrl * im_kp;
            im_y_kappa_integrand = re_y_alpha_itgrl * im_kp + im_y_alpha_itgrl * re_kp;
            re_z_kappa_integrand = re_z_alpha_itgrl * re_kp - im_z_alpha_itgrl * im_kp;
            im_z_kappa_integrand = re_z_alpha_itgrl * im_kp + im_z_alpha_itgrl * re_kp;

            if (i_k>0)
            {
                re_x_result += 0.5 * (re_d_kp * (re_x_kappa_integrand + re_x_kappa_integrand_old)
                                      - im_d_kp * (im_x_kappa_integrand + im_x_kappa_integrand_old));
                im_x_result += 0.5 * (re_d_kp * (im_x_kappa_integrand + im_x_kappa_integrand_old)
                                      + im_d_kp * (re_x_kappa_integrand + re_x_kappa_integrand_old));
                re_y_result += 0.5 * (re_d_kp * (re_y_kappa_integrand + re_y_kappa_integrand_old)
                                      - im_d_kp * (im_y_kappa_integrand + im_y_kappa_integrand_old));
                im_y_result += 0.5 * (re_d_kp * (im_y_kappa_integrand + im_y_kappa_integrand_old)
                                      + im_d_kp * (re_y_kappa_integrand + re_y_kappa_integrand_old));
                re_z_result += 0.5 * (re_d_kp * (re_z_kappa_integrand + re_z_kappa_integrand_old)
                                      - im_d_kp * (im_z_kappa_integrand + im_z_kappa_integrand_old));
                im_z_result += 0.5 * (re_d_kp * (im_z_kappa_integrand + im_z_kappa_integrand_old)
                                      + im_d_kp * (re_z_kappa_integrand + re_z_kappa_integrand_old));
            }
        }
        re_e_x[i] = re_x_result;
        im_e_x[i] = im_x_result;
        re_e_y[i] = re_y_result;
        im_e_y[i] = im_y_result;
        re_e_z[i] = re_z_result;
        im_e_z[i] = im_z_result;
    }
"""


#https://www.grammarbot.io/quickstart
from grammarbot import GrammarBotClient
from numpy.random import choice

# Creating the client
# ===================
client = GrammarBotClient()

#API key is optional
client = GrammarBotClient(api_key='KS9C5N3Y')

# ===================
RANDOM_CHANCE = 0.6

# Generates weights for use in numpy's random choice
# Output looks somewhat like this: [0.41957572502685286, 0.2517454350161117, 0.15104726100966703, 0.09062835660580021, 0.05437701396348013, 0.03262620837808808 ...]
# The output will always sum to 1
def generateWeights(size):
    weights = [RANDOM_CHANCE]
    for i in range(size-1):
        weights.append(weights[i]*RANDOM_CHANCE)
    normalized = [float(i)/sum(weights) for i in weights]
    return normalized

# Calls the grammarbot api to correct text
def grammarCheck(text):
    if text is None or len(text) < 3:
        text = ""
        return ""
    #print(str(text))
    res = client.check(str(text))
    changedLength = 0
    for match in res.matches:
        offset = match.replacement_offset + changedLength #get corrected offset
        if match.replacements is not None and len(match.replacements) > 0:
            replacement = choice(match.replacements,  size=1, p=generateWeights(len(match.replacements)), replace=False)[0] #get a random correction
            text = text.replace(text[offset:offset+match.replacement_length],replacement,1) #replace whatever's there with the correction
            changedLength += len(str(replacement)) - match.replacement_length #track changes to string length
    return text

# Recursively grammar check multiple times
# Ideal for long text where certain corrections could cause issues
def grammarCheckWithPasses(text, passes):
    if(int(passes) == 0):
        return text
    else:
        try:
            return grammarCheckWithPasses(grammarCheck(text), int(passes)-1)
        except:
            return text

'''
text = 'By mocking basic human nature, Twain\'s legendary commentary retains its significance deep into the 21st century. The overall point of the book, that minorities are just as human as everyone else and deserve just as many rights as everyone else, is still a 21st century goal. However, possibly Twain\'s most pertinent criticism is the irony of Pap\'s "You call this a government" speech. In it, Pap argues that the government is so outrageous, he might as well not vote. In the trying times of 2016 and the upcoming 2020 election, this is the exact mindset of many Americans. How has Twain predicted so far ahead? Through a cunning analysis of human nature: while he couldn\'t have known the names of the candidates, he knew the mindset of the country and the commonfolk which form it. While his caricature of Pap is obviously wrong, we start debating when we talk about modern politics, even though the paralells are so clear. This was Twain\'s true genious and the reason behind writing his novels.'
text = grammarCheck(text)
print(text)
'''


import argparse
import json
import math
import os
from datetime import datetime
from collections import OrderedDict
import numpy as np
import torch
import torchvision.transforms as T
import torch.nn.functional as F
from imageio import imwrite

parser = argparse.ArgumentParser()
parser.add_argument('--model_path', required=True)
parser.add_argument('--output_dir', default='outputs')
args = parser.parse_args()

def get_model():
    if not os.path.isdir(os.path.join('gui', args.output_dir)):
        print('Output directory "%s" does not exist; creating it' % args.output_dir)
        os.makedirs(os.path.join('gui', args.output_dir))

    # device = torch.device('cuda:0')

    # Load the model, with a bit of care in case there are no GPUs
    model = torch.jit.load(args.model_path)
    model.colors = torch.FloatTensor(np.load('./color.npy')).cuda()
    # model.eval()
    # model.to(device)

    return model

def json_to_img(scene_graph, model):
    output_dir = args.output_dir
    scene_graphs = json_to_scene_graph(scene_graph)
    current_time = datetime.now().strftime('%b%d_%H-%M-%S')

    # Run the model forward
    with open('./data/coco_vocab.json', 'r') as f:
        vocab = json.load(f)
    bbox_lt, attributes, label = encode_scene_graphs(scene_graphs, vocab)
    with torch.no_grad():
        z_obj = torch.from_numpy(truncted_random(num_o=8, thres=2.0)).float().cuda()
        z_im = torch.from_numpy(truncted_random(num_o=1, thres=2.0)).view(1, -1).float().cuda()
        # imgs, bbox_wh, panoptic_bbox = model(z_obj, bbox_lt=bbox_lt, attr=attributes, y=label.squeeze(dim=-1))
        imgs, bbox_wh, layout = model(z_obj, bbox_lt[:, :, :2], attributes, z_im, label.squeeze(dim=-1))
        # imgs_org, bmask_org = org_model(z_obj, torch.cat([bbox_lt, bbox_wh], 2), y=label.squeeze(dim=-1))
    imgs = imagenet_deprocess_batch(imgs)
    # imgs_org = imagenet_deprocess_batch(imgs_org)

    # Save the generated image
    for i in range(imgs.shape[0]):
        img_np = imgs[i].numpy().transpose(1, 2, 0).astype('uint8')
        return_img_path = os.path.join(output_dir, 'img{}.png'.format(current_time))
        img_path = os.path.join('gui', output_dir, 'img{}.png'.format(current_time))
        imwrite(img_path, img_np)

        # img_np_org = imgs_org[i].numpy().transpose(1, 2, 0).astype('uint8')
        # return_img_path_org = os.path.join('images', output_dir, 'img_org{}.png'.format(current_time))
        # img_path_org = os.path.join('scripts', 'gui', return_img_path_org)
        # imwrite(img_path_org, img_np_org)

    # Save the generated layouts image
    for i in range(imgs.shape[0]):
        # b, h, w = bmask_org.size(0), bmask_org.size(2), bmask_org.size(3)
        # bmask_org = torch.cat([torch.zeros([b, 1, h, w], device=bmask_org.device).float(), bmask_org], 1)
        # bmask_org = torch.argmax(bmask_org, 1)
        # num_class = 9
        # one_hot = F.one_hot(bmask_org, num_class).permute(0, 3, 1, 2).float()
        # one_hot[:, 0] = 0
        # img_layout_np = one_hot_to_rgb(one_hot, model.colors[:num_class, :])[0].numpy().transpose(1, 2, 0).astype(
        #     'uint8')
        # return_img_layout_path = os.path.join('images', output_dir, 'img_layout{}.png'.format(current_time))
        # img_layout_path = os.path.join('scripts', 'gui', return_img_layout_path)
        # # vis.add_boxes_to_layout(img_layout_np, scene_graphs[i]['objects'], boxes_pred, img_layout_path,
        # #                         colors=obj_colors)
        # imwrite(img_layout_path, img_layout_np)

        b, h, w = layout.size(0), layout.size(2), layout.size(3)
        layout = torch.cat([torch.zeros([b, 1, h, w], device=layout.device).float(), layout], 1)
        layout = torch.argmax(layout, 1)
        num_class = 9
        one_hot = F.one_hot(layout, num_class).permute(0, 3, 1, 2).float()
        one_hot[:, 0] = 0
        layout = one_hot_to_rgb(one_hot, model.colors[:num_class, :])[0].numpy().transpose(1, 2, 0).astype('uint8')
        return_layout_path = os.path.join(output_dir, 'layouts{}.png'.format(current_time))
        layout_path = os.path.join('gui', output_dir, 'layouts{}.png'.format(current_time))
        imwrite(layout_path, layout)

    return return_img_path, return_layout_path

def truncted_random(num_o=8, thres=1.0, bs=1):
    z = np.ones((1, num_o, 128)) * 100
    for i in range(num_o):
        for j in range(128):
            while z[0, i, j] > thres or z[0, i, j] < - thres:
                z[0, i, j] = np.random.normal()
    z = np.tile(z, (bs, 1, 1))
    return z

def one_hot_to_rgb(one_hot, colors):
    one_hot_3d = torch.einsum('abcd,be->aecd', (one_hot.cpu(), colors.cpu()))
    one_hot_3d *= (255.0 / one_hot_3d.max())
    return one_hot_3d

def encode_scene_graphs(scene_graphs, vocab, num_objs=8):
    """
    Encode one or more scene graphs using this model's vocabulary. Inputs to
    this method are scene graphs represented as dictionaries like the following:
    {
      "objects": ["cat", "dog", "sky"],
      "relationships": [
        [0, "next to", 1],
        [0, "beneath", 2],
        [2, "above", 1],
      ]
    }
    This scene graph has three relationshps: cat next to dog, cat beneath sky,
    and sky above dog.
    Inputs:
    - scene_graphs: A dictionary giving a single scene graph, or a list of
      dictionaries giving a sequence of scene graphs.
    Returns a tuple of LongTensors (objs, triples, obj_to_img) that have the
    same semantics as self.forward. The returned LongTensors will be on the
    same device as the model parameters.
    """
    if isinstance(scene_graphs, dict):
        # We just got a single scene graph, so promote it to a list
        scene_graphs = [scene_graphs]

    objs, triples, obj_to_img = [], [], []
    all_attributes = []
    all_features = []
    obj_offset = 0
    bbox_lt = []
    bbox_center = []
    print(scene_graphs)
    for i, sg in enumerate(scene_graphs):
        attributes = torch.zeros([num_objs, 25 + 10], dtype=torch.float).cuda()
        # Insert dummy __image__ object and __in_image__ relationships
        boxes = sg['boxes']
        image_idx = len(sg['objects'])

        for bbox in boxes:
            sx0, sy0, sx1, sy1 = bbox
            bbox_lt.append([sx0, sy0])
            bbox_center.append([(sx0+sx1)/2, (sy0+sy1)/2])
        for obj in sg['objects']:
            obj_idx = vocab['object_name_to_idx'][obj]
            objs.append(obj_idx)
            obj_to_img.append(i)
        for i, size_attr in enumerate(sg['attributes']['size']):
            attributes[i, size_attr] = 1
        for i, location_attr in enumerate(sg['attributes']['location']):
            attributes[i, location_attr + 10] = 1

        for _ in range(len(boxes), num_objs):
            bbox_lt.append([-0.6, -0.6])
            bbox_center.append([-0.6, -0.6])
            objs.append(0)
        all_attributes.append(attributes.unsqueeze(0))

    objs = torch.tensor(objs, dtype=torch.int64).unsqueeze(0).cuda()
    attributes = torch.cat(all_attributes, 0)
    bbox_lt = torch.FloatTensor(bbox_lt).unsqueeze(0).cuda()
    bbox_center = torch.FloatTensor(bbox_center).unsqueeze(0).cuda()

    return bbox_lt, attributes, objs

def json_to_scene_graph(json_text):
    scene = json.loads(json_text)
    if len(scene) == 0:
        return []
    image_id = scene['image_id']
    scene = scene['objects']
    objects = [i['text'] for i in scene]
    relationships = []
    size = []
    location = []
    features = []
    bbox_lt = []
    attr = []
    boxes = []
    for i in range(0, len(objects)):
        obj_s = scene[i]
        # Check for inside / surrounding

        sx0 = obj_s['left']
        sy0 = obj_s['top']
        sx1 = obj_s['width'] + sx0
        sy1 = obj_s['height'] + sy0
        bbox_lt.append([sx0, sy0])
        boxes.append([sx0, sy0, sx1, sy1])

        margin = (obj_s['size'] + 1) / 10 / 2
        mean_x_s = 0.5 * (sx0 + sx1)
        mean_y_s = 0.5 * (sy0 + sy1)

        sx0 = max(0, mean_x_s - margin)
        sx1 = min(1, mean_x_s + margin)
        sy0 = max(0, mean_y_s - margin)
        sy1 = min(1, mean_y_s + margin)

        size.append(obj_s['size'])
        location.append(obj_s['location'])

        features.append(obj_s['feature'])
        if i == len(objects) - 1:
            continue

        obj_o = scene[i + 1]
        ox0 = obj_o['left']
        oy0 = obj_o['top']
        ox1 = obj_o['width'] + ox0
        oy1 = obj_o['height'] + oy0

        mean_x_o = 0.5 * (ox0 + ox1)
        mean_y_o = 0.5 * (oy0 + oy1)
        d_x = mean_x_s - mean_x_o
        d_y = mean_y_s - mean_y_o
        theta = math.atan2(d_y, d_x)

        margin = (obj_o['size'] + 1) / 10 / 2
        ox0 = max(0, mean_x_o - margin)
        ox1 = min(1, mean_x_o + margin)
        oy0 = max(0, mean_y_o - margin)
        oy1 = min(1, mean_y_o + margin)

        if sx0 < ox0 and sx1 > ox1 and sy0 < oy0 and sy1 > oy1:
            p = 'surrounding'
        elif sx0 > ox0 and sx1 < ox1 and sy0 > oy0 and sy1 < oy1:
            p = 'inside'
        elif theta >= 3 * math.pi / 4 or theta <= -3 * math.pi / 4:
            p = 'left of'
        elif -3 * math.pi / 4 <= theta < -math.pi / 4:
            p = 'above'
        elif -math.pi / 4 <= theta < math.pi / 4:
            p = 'right of'
        elif math.pi / 4 <= theta < 3 * math.pi / 4:
            p = 'below'
        relationships.append([i, p, i + 1])

    return [{'objects': objects, 'relationships': relationships, 'attributes': {'size': size, 'location': location},
             'features': features, 'image_id': image_id, 'boxes': boxes}]

def rescale(x):
    lo, hi = x.min(), x.max()
    return x.sub(lo).div(hi - lo)

def imagenet_deprocess(rescale_image=True, MEAN=[0.5, 0.5, 0.5], STD=[0.5, 0.5, 0.5]):
    INV_MEAN = [-m for m in MEAN]
    INV_STD = [1.0 / s for s in STD]
    transforms = [
        T.Normalize(mean=[0, 0, 0], std=INV_STD),
        T.Normalize(mean=INV_MEAN, std=[1.0, 1.0, 1.0]),
    ]
    if rescale_image:
        transforms.append(rescale)
    return T.Compose(transforms)

def imagenet_deprocess_batch(imgs, rescale=True):
    """
    Input:
    - imgs: FloatTensor of shape (N, C, H, W) giving preprocessed images
    Output:
    - imgs_de: ByteTensor of shape (N, C, H, W) giving deprocessed images
      in the range [0, 255]
    """
    if isinstance(imgs, torch.autograd.Variable):
        imgs = imgs.data
    imgs = imgs.cpu().clone()
    deprocess_fn = imagenet_deprocess(rescale_image=rescale)
    imgs_de = []
    for i in range(imgs.size(0)):
        img_de = deprocess_fn(imgs[i])[None]
        img_de = img_de.mul(255).clamp(0, 255)
        imgs_de.append(img_de)
    imgs_de = torch.cat(imgs_de, dim=0)
    return imgs_de