import os
import sys
import logging
import argparse
import itertools
import torch
import inspect
from os.path import join as pjoin
import numpy as np

#relative path hack
test_path = os.path.join(os.path.dirname(__file__), '..', 'Common')
test = os.path.join(sys.path[0], '..')
currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
python_ssd_traindir = os.path.join(parentdir,'python','training','detection','ssd')
sys.path.append(python_ssd_traindir)

from torch.utils.data import DataLoader, ConcatDataset
from torch.optim.lr_scheduler import CosineAnnealingLR, MultiStepLR

from vision.utils.misc import str2bool, Timer, freeze_net_layers, store_labels
from vision.ssd.ssd import MatchPrior
from vision.ssd.vgg_ssd import create_vgg_ssd
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.nn.multibox_loss import MultiboxLoss
from vision.ssd.config import vgg_ssd_config
from vision.ssd.config import mobilenetv1_ssd_config
from vision.ssd.config import squeezenet_ssd_config
from vision.ssd.data_preprocessing import TrainAugmentation, TestTransform

from vision.ssd.vgg_ssd import create_vgg_ssd, create_vgg_ssd_predictor
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd, create_mobilenetv1_ssd_predictor
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite, create_mobilenetv1_ssd_lite_predictor
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite, create_squeezenet_ssd_lite_predictor
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite, create_mobilenetv2_ssd_lite_predictor
from vision.utils.misc import Timer

import cv2

def main(args):
    net_type = args.net_type

    model_path = args.model_path
    label_path = args.label_path
    image_folder = args.image_folder
    output_folder = args.output_folder

    class_names = [name.strip() for name in open(label_path).readlines()]
    DEVICE = torch.device("cuda:0")

    if net_type == 'vgg16-ssd':
        net = create_vgg_ssd(len(class_names), is_test=True)
    elif net_type == 'mb1-ssd':
        net = create_mobilenetv1_ssd(len(class_names), is_test=True)
    elif net_type == 'mb1-ssd-lite':
        net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
    elif net_type == 'mb2-ssd-lite':
        net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
    elif net_type == 'sq-ssd-lite':
        net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
    else:
        print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
        sys.exit(1)
    net.load(model_path)

    if net_type == 'vgg16-ssd':
        predictor = create_vgg_ssd_predictor(net, candidate_size=200,device=DEVICE)
    elif net_type == 'mb1-ssd':
        predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200,device=DEVICE)
    elif net_type == 'mb1-ssd-lite':
        predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200,device=DEVICE)
    elif net_type == 'mb2-ssd-lite':
        predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200,device=DEVICE)
    elif net_type == 'sq-ssd-lite':
        predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200,device=DEVICE)
    else:
        predictor = create_vgg_ssd_predictor(net, candidate_size=200,device=DEVICE)

    os.makedirs(output_folder,exist_ok=True)

    for file in os.listdir(image_folder):
        if os.path.splitext(file)[-1] != '.jpg':
            continue
        image_path = pjoin(image_folder,file)
        orig_image = cv2.imread(image_path)
        image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)

        # image = image.to(DEVICE)
        boxes, labels, probs = predictor.predict(image, 10, 0.4)
        if boxes.size(0) == 0:
            print('no anom found, skipping')
            continue
        for i in range(boxes.size(0)):
            box = boxes[i, :]
            cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), (255, 255, 0), 4)
            #label = f"""{voc_dataset.class_names[labels[i]]}: {probs[i]:.2f}"""
            label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
            cv2.putText(orig_image, label,
                        (box[0] + 20, box[1] + 40),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        1,  # font scale
                        (255, 0, 255),
                        2)  # line type

        cv2.imwrite(pjoin(output_folder,file), orig_image)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description='')

    # Params for datasets
    parser.add_argument("--net_type", required=True,
                        dest='net_type',
                        type=str,
                        help='')
    parser.add_argument("--model_path", required=True,
                        dest='model_path',
                        type=str,
                        help='')
    parser.add_argument("--label_path", required=True,
                        dest='label_path',
                        type=str,
                        help='')
    parser.add_argument("--image_folder", required=True,
                        dest='image_folder',
                        type=str,
                        help='')
    parser.add_argument("--output_folder", required=True,
                        dest='output_folder',
                        type=str,
                        help='')

    args = parser.parse_args()
    main(args)


# Copyright 2019 Xilinx Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import os
import cv2
import numpy as np
import tensorflow as tf

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(BASE_DIR + '/test')

# import model_utils.preprocess as preprocess
from model_utils.preprocess import PREPROCESS_FUNC
from model_utils.config import CONFIG_MAP

def get_config(key=None, default_value=None):
  if not key:
    raise ValueError("Please assign a key.")
  if not default_value:
    raise ValueEror("Please assign a default_value")

  config = os.environ
  if key in config:
    value = config[key]
    print("Get {} from env: {}".format(key, value))
    return value
  else:
    print("Fail to get {} from env, use default value {}".format(
        key, default_value))
    return default_value

model_type = get_config(key="MODEL_TYPE", default_value="ssdlite_mobilenet_v2")
calib_image_dir = get_config(
    key="CALIB_IMAGE_DIR",
    default_value=
    "data/calibration_data/coco_images/")
calib_image_list = get_config(
    key="CALIB_IMAGE_LIST",
    default_value="data/calib_list.txt")
calib_batch_size = int(get_config(key="CALIB_BATCH_SIZE", default_value=1))
input_height = int(get_config(key="INPUT_HEIGHT", default_value=300))
input_width = int(get_config(key="INPUT_WIDTH", default_value=300))

def calib_input(iter):
  images = []
  line = open(calib_image_list).readlines()
  with tf.Graph().as_default():
    raw_image = tf.placeholder(tf.float32,
                               shape=(None, None, None, 3),
                               name="raw_image")
    # preprocess_func = preprocess.PREPROCESS_FUNC[model_type]
    preprocess_func = PREPROCESS_FUNC[CONFIG_MAP[model_type].feature_extractor_type]
    preprocessed_image = preprocess_func(raw_image, input_height, input_width)
    sess = tf.Session()
    for index in range(0, calib_batch_size):
      curline = line[iter * calib_batch_size + index]
      calib_image_name = curline.strip()
      image_path = os.path.join(calib_image_dir, calib_image_name + ".jpg")
      image = cv2.imread(image_path)
      height, width = image.shape[0:2]
      image = image[:, :, ::-1]  # BGR to RGB
      image = np.array(image, dtype=np.float32)
      image = np.expand_dims(image, 0)
      image_tensor = tf.get_default_graph().get_tensor_by_name("raw_image:0")
      pre_image = sess.run(preprocessed_image, feed_dict={image_tensor: image})
      pre_image = np.squeeze(pre_image)
      images.append(pre_image)
    sess.close()
  return {"image_tensor": images}


# coding: utf-8

# # Language Translation
# In this project, you’re going to take a peek into the realm of neural network machine translation.  You’ll be training a sequence to sequence model on a dataset of English and French sentences that can translate new sentences from English to French.
# ## Get the Data
# Since translating the whole language of English to French will take lots of time to train, we have provided you with a small portion of the English corpus.

# In[4]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
import helper
import problem_unittests as tests

source_path = 'data/small_vocab_en'
target_path = 'data/small_vocab_fr'
source_text = helper.load_data(source_path)
target_text = helper.load_data(target_path)

# ## Explore the Data
# Play around with view_sentence_range to view different parts of the data.

# In[5]:

view_sentence_range = (0, 10)

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
import numpy as np

print('Dataset Stats')
print('Roughly the number of unique words: {}'.format(len({word: None for word in source_text.split()})))

sentences = source_text.split('\n')
word_counts = [len(sentence.split()) for sentence in sentences]
print('Number of sentences: {}'.format(len(sentences)))
print('Average number of words in a sentence: {}'.format(np.average(word_counts)))

print()
print('English sentences {} to {}:'.format(*view_sentence_range))
print('\n'.join(source_text.split('\n')[view_sentence_range[0]:view_sentence_range[1]]))
print()
print('French sentences {} to {}:'.format(*view_sentence_range))
print('\n'.join(target_text.split('\n')[view_sentence_range[0]:view_sentence_range[1]]))

# ## Implement Preprocessing Function
# ### Text to Word Ids
# As you did with other RNNs, you must turn the text into a number so the computer can understand it. In the function `text_to_ids()`, you'll turn `source_text` and `target_text` from words to ids.  However, you need to add the `<EOS>` word id at the end of each sentence from `target_text`.  This will help the neural network predict when the sentence should end.
#
# You can get the `<EOS>` word id by doing:
# ```python
# target_vocab_to_int['<EOS>']
# ```
# You can get other word ids using `source_vocab_to_int` and `target_vocab_to_int`.

# In[6]:

def text_to_ids(source_text, target_text, source_vocab_to_int, target_vocab_to_int):
    """
    Convert source and target text to proper word ids
    :param source_text: String that contains all the source text.
    :param target_text: String that contains all the target text.
    :param source_vocab_to_int: Dictionary to go from the source words to an id
    :param target_vocab_to_int: Dictionary to go from the target words to an id
    :return: A tuple of lists (source_id_text, target_id_text)
    """
    # TODO: Implement Function
    source_id_text = [[source_vocab_to_int[word] for word in sentence.split()] for sentence in source_text.split('\n')]

    eos_word = target_vocab_to_int['<EOS>']
    target_id_text = [[target_vocab_to_int[word] for word in sentence.split()] + [eos_word] for sentence in target_text.split('\n')]

    #print(target_id_text)
    return (source_id_text, target_id_text)

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_text_to_ids(text_to_ids)

# ### Preprocess all the data and save it
# Running the code cell below will preprocess all the data and save it to file.

# In[7]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
helper.preprocess_and_save_data(source_path, target_path, text_to_ids)

# # Check Point
# This is your first checkpoint. If you ever decide to come back to this notebook or have to restart the notebook, you can start from here. The preprocessed data has been saved to disk.

# In[8]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
import numpy as np
import helper

(source_int_text, target_int_text), (source_vocab_to_int, target_vocab_to_int), _ = helper.load_preprocess()

# ### Check the Version of TensorFlow and Access to GPU
# This will check to make sure you have the correct version of TensorFlow and access to a GPU

# In[9]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
from distutils.version import LooseVersion
import warnings
import tensorflow as tf

# Check TensorFlow Version
assert LooseVersion(tf.__version__) in [LooseVersion('1.0.0'), LooseVersion('1.0.1')], 'This project requires TensorFlow version 1.0  You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
    warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))

# ## Build the Neural Network
# You'll build the components necessary to build a Sequence-to-Sequence model by implementing the following functions below:
# - `model_inputs`
# - `process_decoding_input`
# - `encoding_layer`
# - `decoding_layer_train`
# - `decoding_layer_infer`
# - `decoding_layer`
# - `seq2seq_model`
#
# ### Input
# Implement the `model_inputs()` function to create TF Placeholders for the Neural Network. It should create the following placeholders:
#
# - Input text placeholder named "input" using the TF Placeholder name parameter with rank 2.
# - Targets placeholder with rank 2.
# - Learning rate placeholder with rank 0.
# - Keep probability placeholder named "keep_prob" using the TF Placeholder name parameter with rank 0.
#
# Return the placeholders in the following the tuple (Input, Targets, Learing Rate, Keep Probability)

# In[10]:

def model_inputs():
    """
    Create TF Placeholders for input, targets, and learning rate.
    :return: Tuple (input, targets, learning rate, keep probability)
    """
    # TODO: Implement Function
    data_input = tf.placeholder(tf.int32, shape=(None, None), name='input')
    data_target = tf.placeholder(tf.int32, shape=(None, None), name='targets')
    learn_rate = tf.placeholder(tf.float32, name='learning_rate')
    keep_prob = tf.placeholder(tf.float32, name='keep_prob')

    return (data_input, data_target, learn_rate, keep_prob)

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_model_inputs(model_inputs)

# ### Process Decoding Input
# Implement `process_decoding_input` using TensorFlow to remove the last word id from each batch in `target_data` and concat the GO ID to the beginning of each batch.

# In[11]:

def process_decoding_input(target_data, target_vocab_to_int, batch_size):
    """
    Preprocess target data for decoding
    :param target_data: Target Placeholder
    :param target_vocab_to_int: Dictionary to go from the target words to an id
    :param batch_size: Batch Size
    :return: Preprocessed target data
    """
    # TODO: Implement Function
    # This is the code from the lesson, any resources tohelp with further understanding of this?
    ending = tf.strided_slice(target_data, [0, 0], [batch_size, -1], [1, 1])
    dec_input = tf.concat([tf.fill([batch_size, 1], target_vocab_to_int['<GO>']), ending], 1)
    return dec_input

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_process_decoding_input(process_decoding_input)

# ### Encoding
# Implement `encoding_layer()` to create a Encoder RNN layer using [`tf.nn.dynamic_rnn()`](https://www.tensorflow.org/api_docs/python/tf/nn/dynamic_rnn).

# In[12]:

def encoding_layer(rnn_inputs, rnn_size, num_layers, keep_prob):
    """
    Create encoding layer
    :param rnn_inputs: Inputs for the RNN
    :param rnn_size: RNN Size
    :param num_layers: Number of layers
    :param keep_prob: Dropout keep probability
    :return: RNN state
    """
    # TODO: Implement Function
    # Initial state of the LSTM memory.
    stacked_lstm_cells = tf.contrib.rnn.MultiRNNCell([build_rnn_cell(rnn_size, keep_prob) for _ in range(num_layers)], state_is_tuple=True)
    outputs, cells_bw_fw = tf.nn.dynamic_rnn(stacked_lstm_cells, rnn_inputs, dtype=tf.float32)

    return cells_bw_fw

# use Tensorflow 1.1 and this olves an error as you will see below
def build_rnn_cell(rnn_size,keep_prob):
    lstm_cell = tf.contrib.rnn.BasicLSTMCell(rnn_size, state_is_tuple=True)
    # Very usefull
    drop = tf.contrib.rnn.DropoutWrapper(lstm_cell, input_keep_prob=1.0, output_keep_prob=keep_prob)
    return drop

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_encoding_layer(encoding_layer)

# ### Decoding - Training
# Create training logits using [`tf.contrib.seq2seq.simple_decoder_fn_train()`](https://www.tensorflow.org/versions/r1.0/api_docs/python/tf/contrib/seq2seq/simple_decoder_fn_train) and [`tf.contrib.seq2seq.dynamic_rnn_decoder()`](https://www.tensorflow.org/versions/r1.0/api_docs/python/tf/contrib/seq2seq/dynamic_rnn_decoder).  Apply the `output_fn` to the [`tf.contrib.seq2seq.dynamic_rnn_decoder()`](https://www.tensorflow.org/versions/r1.0/api_docs/python/tf/contrib/seq2seq/dynamic_rnn_decoder) outputs.

# In[13]:

def decoding_layer_train(encoder_state, dec_cell, dec_embed_input, sequence_length, decoding_scope, output_fn, keep_prob):
    """
    Create a decoding layer for training
    :param encoder_state: Encoder State
    :param dec_cell: Decoder RNN Cell
    :param dec_embed_input: Decoder embedded input
    :param sequence_length: Sequence Length
    :param decoding_scope: TensorFlow Variable Scope for decoding
    :param output_fn: Function to apply the output layer
    :param keep_prob: Dropout keep probability
    :return: Train Logits
    """
    # TODO: Implement Function
    train_dynamic_fn = tf.contrib.seq2seq.simple_decoder_fn_train(encoder_state)
    train_outputs, train_final_state, train_final_ctx_state = tf.contrib.seq2seq.dynamic_rnn_decoder(cell=dec_cell, decoder_fn=train_dynamic_fn, inputs=dec_embed_input, sequence_length=sequence_length, scope=decoding_scope)
    train_logits = output_fn(train_outputs)

    return train_logits

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_decoding_layer_train(decoding_layer_train)

# ### Decoding - Inference
# Create inference logits using [`tf.contrib.seq2seq.simple_decoder_fn_inference()`](https://www.tensorflow.org/versions/r1.0/api_docs/python/tf/contrib/seq2seq/simple_decoder_fn_inference) and [`tf.contrib.seq2seq.dynamic_rnn_decoder()`](https://www.tensorflow.org/versions/r1.0/api_docs/python/tf/contrib/seq2seq/dynamic_rnn_decoder).

# In[14]:

def decoding_layer_infer(encoder_state, dec_cell, dec_embeddings, start_of_sequence_id, end_of_sequence_id,
                         maximum_length, vocab_size, decoding_scope, output_fn, keep_prob):
    """
    Create a decoding layer for inference
    :param encoder_state: Encoder state
    :param dec_cell: Decoder RNN Cell
    :param dec_embeddings: Decoder embeddings
    :param start_of_sequence_id: GO ID
    :param end_of_sequence_id: EOS Id
    :param maximum_length: The maximum allowed time steps to decode
    :param vocab_size: Size of vocabulary
    :param decoding_scope: TensorFlow Variable Scope for decoding
    :param output_fn: Function to apply the output layer
    :param keep_prob: Dropout keep probability
    :return: Inference Logits
    """
    # TODO: Implement Function
    infer_dynamic_fn = tf.contrib.seq2seq.simple_decoder_fn_inference(output_fn, encoder_state,
                                                                      dec_embeddings, start_of_sequence_id,
                                                                      end_of_sequence_id, maximum_length, vocab_size)
    infer_outputs, infer_final_state, infer_final_ctx_state = tf.contrib.seq2seq.dynamic_rnn_decoder(cell=dec_cell,
                                                                                decoder_fn=infer_dynamic_fn,
                                                                                scope=decoding_scope)
    infer_logits = infer_outputs
    return infer_logits

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_decoding_layer_infer(decoding_layer_infer)

# ### Build the Decoding Layer
# Implement `decoding_layer()` to create a Decoder RNN layer.
#
# - Create RNN cell for decoding using `rnn_size` and `num_layers`.
# - Create the output fuction using [`lambda`](https://docs.python.org/3/tutorial/controlflow.html#lambda-expressions) to transform it's input, logits, to class logits.
# - Use the your `decoding_layer_train(encoder_state, dec_cell, dec_embed_input, sequence_length, decoding_scope, output_fn, keep_prob)` function to get the training logits.
# - Use your `decoding_layer_infer(encoder_state, dec_cell, dec_embeddings, start_of_sequence_id, end_of_sequence_id, maximum_length, vocab_size, decoding_scope, output_fn, keep_prob)` function to get the inference logits.
#
# Note: You'll need to use [tf.variable_scope](https://www.tensorflow.org/api_docs/python/tf/variable_scope) to share variables between training and inference.

# In[15]:

def decoding_layer(dec_embed_input, dec_embeddings, encoder_state, vocab_size, sequence_length, rnn_size,
                   num_layers, target_vocab_to_int, keep_prob):
    """
    Create decoding layer
    :param dec_embed_input: Decoder embedded input
    :param dec_embeddings: Decoder embeddings
    :param encoder_state: The encoded state
    :param vocab_size: Size of vocabulary
    :param sequence_length: Sequence Length
    :param rnn_size: RNN Size
    :param num_layers: Number of layers
    :param target_vocab_to_int: Dictionary to go from the target words to an id
    :param keep_prob: Dropout keep probability
    :return: Tuple of (Training Logits, Inference Logits)
    """
    # TODO: Implement Function
    # Initial state of the LSTM memory.
    dec_stacked_cells = tf.contrib.rnn.MultiRNNCell([build_rnn_cell(rnn_size, keep_prob) for _ in range(num_layers)], state_is_tuple=True)
    end_seq_id = target_vocab_to_int['<EOS>']
    start_seq_id = target_vocab_to_int['<GO>']

    with tf.variable_scope("decoding", reuse=None) as decoding_scope:
        # as shown in problem_unittests.py
        output_fn = lambda x: tf.contrib.layers.fully_connected(x, vocab_size, None, scope=decoding_scope)
    with tf.variable_scope("decoding", reuse=None) as decoding_scope:
        train_logits = decoding_layer_train(encoder_state, dec_stacked_cells,
                                            dec_embed_input, sequence_length, decoding_scope,
                                            output_fn, keep_prob)
    # Error when reuse=True
    with tf.variable_scope("decoding", reuse=True) as decoding_scope:
        infer_logits = decoding_layer_infer(encoder_state, dec_stacked_cells, dec_embeddings,
                                            start_seq_id, end_seq_id, sequence_length-1, vocab_size,
                                            decoding_scope, output_fn, keep_prob)

    return (train_logits, infer_logits)

# use Tensorflow 1.1 and this olves an error as you will see below
def build_rnn_cell(rnn_size,keep_prob):
    lstm_cell = tf.contrib.rnn.BasicLSTMCell(rnn_size, state_is_tuple=True)
    # Very usefull
    drop = tf.contrib.rnn.DropoutWrapper(lstm_cell, input_keep_prob=1.0, output_keep_prob=keep_prob)
    return drop

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_decoding_layer(decoding_layer)

# ### Build the Neural Network
# Apply the functions you implemented above to:
#
# - Apply embedding to the input data for the encoder.
# - Encode the input using your `encoding_layer(rnn_inputs, rnn_size, num_layers, keep_prob)`.
# - Process target data using your `process_decoding_input(target_data, target_vocab_to_int, batch_size)` function.
# - Apply embedding to the target data for the decoder.
# - Decode the encoded input using your `decoding_layer(dec_embed_input, dec_embeddings, encoder_state, vocab_size, sequence_length, rnn_size, num_layers, target_vocab_to_int, keep_prob)`.

# In[16]:

def seq2seq_model(input_data, target_data, keep_prob, batch_size, sequence_length, source_vocab_size, target_vocab_size,
                  enc_embedding_size, dec_embedding_size, rnn_size, num_layers, target_vocab_to_int):
    """
    Build the Sequence-to-Sequence part of the neural network
    :param input_data: Input placeholder
    :param target_data: Target placeholder
    :param keep_prob: Dropout keep probability placeholder
    :param batch_size: Batch Size
    :param sequence_length: Sequence Length
    :param source_vocab_size: Source vocabulary size
    :param target_vocab_size: Target vocabulary size
    :param enc_embedding_size: Decoder embedding size
    :param dec_embedding_size: Encoder embedding size
    :param rnn_size: RNN Size
    :param num_layers: Number of layers
    :param target_vocab_to_int: Dictionary to go from the target words to an id
    :return: Tuple of (Training Logits, Inference Logits)
    """
    # TODO: Implement Function
    #1
    rnn_inputs = tf.contrib.layers.embed_sequence(input_data, source_vocab_size, enc_embedding_size)
    #2
    input_encoded_state = encoding_layer(rnn_inputs, rnn_size, num_layers, keep_prob)
    #3
    target_output = process_decoding_input(target_data, target_vocab_to_int, batch_size)
    #4
    dec_embeds = tf.Variable(tf.random_uniform([target_vocab_size, dec_embedding_size]))
    dec_embeds_input = tf.nn.embedding_lookup(dec_embeds, target_output)
    #5
    train_logits, infer_logits = decoding_layer(dec_embeds_input, dec_embeds, input_encoded_state, target_vocab_size,
                                                sequence_length, rnn_size, num_layers, target_vocab_to_int, keep_prob)
    return (train_logits, infer_logits)

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_seq2seq_model(seq2seq_model)

# ## Neural Network Training
# ### Hyperparameters
# Tune the following parameters:
#
# - Set `epochs` to the number of epochs.
# - Set `batch_size` to the batch size.
# - Set `rnn_size` to the size of the RNNs.
# - Set `num_layers` to the number of layers.
# - Set `encoding_embedding_size` to the size of the embedding for the encoder.
# - Set `decoding_embedding_size` to the size of the embedding for the decoder.
# - Set `learning_rate` to the learning rate.
# - Set `keep_probability` to the Dropout keep probability

# In[28]:

# Number of Epochs
epochs = 6
# Batch Size
batch_size = 128
# RNN Size
rnn_size = 128
# Number of Layers
num_layers = 1
# Embedding Size
encoding_embedding_size = 256
decoding_embedding_size = 256
# Learning Rate
learning_rate = 0.007
# Dropout Keep Probability
keep_probability = 0.5

# ### Build the Graph
# Build the graph using the neural network you implemented.

# In[29]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
save_path = 'checkpoints/dev'
(source_int_text, target_int_text), (source_vocab_to_int, target_vocab_to_int), _ = helper.load_preprocess()
max_source_sentence_length = max([len(sentence) for sentence in source_int_text])

train_graph = tf.Graph()
with train_graph.as_default():
    input_data, targets, lr, keep_prob = model_inputs()
    sequence_length = tf.placeholder_with_default(max_source_sentence_length, None, name='sequence_length')
    input_shape = tf.shape(input_data)

    train_logits, inference_logits = seq2seq_model(
        tf.reverse(input_data, [-1]), targets, keep_prob, batch_size, sequence_length, len(source_vocab_to_int), len(target_vocab_to_int),
        encoding_embedding_size, decoding_embedding_size, rnn_size, num_layers, target_vocab_to_int)

    tf.identity(inference_logits, 'logits')
    with tf.name_scope("optimization"):
        # Loss function
        cost = tf.contrib.seq2seq.sequence_loss(
            train_logits,
            targets,
            tf.ones([input_shape[0], sequence_length]))

        # Optimizer
        optimizer = tf.train.AdamOptimizer(lr)

        # Gradient Clipping
        gradients = optimizer.compute_gradients(cost)
        capped_gradients = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in gradients if grad is not None]
        train_op = optimizer.apply_gradients(capped_gradients)

# ### Train
# Train the neural network on the preprocessed data. If you have a hard time getting a good loss, check the forums to see if anyone is having the same problem.

# In[30]:

import matplotlib.pyplot as plt

# In[31]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
import time

def get_accuracy(target, logits):
    """
    Calculate accuracy
    """
    max_seq = max(target.shape[1], logits.shape[1])
    if max_seq - target.shape[1]:
        target = np.pad(
            target,
            [(0,0),(0,max_seq - target.shape[1])],
            'constant')
    if max_seq - logits.shape[1]:
        logits = np.pad(
            logits,
            [(0,0),(0,max_seq - logits.shape[1]), (0,0)],
            'constant')

    return np.mean(np.equal(target, np.argmax(logits, 2)))

train_source = source_int_text[batch_size:]
train_target = target_int_text[batch_size:]

valid_source = helper.pad_sentence_batch(source_int_text[:batch_size])
valid_target = helper.pad_sentence_batch(target_int_text[:batch_size])

with tf.Session(graph=train_graph) as sess:
    sess.run(tf.global_variables_initializer())

    losses = {'train':[], 'validation':[]}
    for epoch_i in range(epochs):
        for batch_i, (source_batch, target_batch) in enumerate(
                helper.batch_data(train_source, train_target, batch_size)):
            start_time = time.time()

            _, loss = sess.run(
                [train_op, cost],
                {input_data: source_batch,
                 targets: target_batch,
                 lr: learning_rate,
                 sequence_length: target_batch.shape[1],
                 keep_prob: keep_probability})

            batch_train_logits = sess.run(
                inference_logits,
                {input_data: source_batch, keep_prob: 1.0})
            batch_valid_logits = sess.run(
                inference_logits,
                {input_data: valid_source, keep_prob: 1.0})

            train_acc = get_accuracy(target_batch, batch_train_logits)
            valid_acc = get_accuracy(np.array(valid_target), batch_valid_logits)
            end_time = time.time()
            #print('Epoch {:>3} Batch {:>4}/{} - Train Accuracy: {:>6.3f}, Validation Accuracy: {:>6.3f}, Loss: {:>6.3f}'.format(epoch_i, batch_i, len(source_int_text) // batch_size, train_acc, valid_acc, loss))
            sys.stdout.write('Epoch {:>3} Batch {:>4}/{} - Train Accuracy: {:>6.3f}, Validation Accuracy: {:>6.3f}, Loss: {:>6.3f}'
                  .format(epoch_i, batch_i, len(source_int_text) // batch_size, train_acc, valid_acc, loss))
            losses['train'].append(train_acc)
            losses['validation'].append(valid_acc)
    print("")
    #plot accuracy
    plt.plot(losses['train'], label='Training loss')
    plt.plot(losses['validation'], label='Validation loss')
    plt.legend()
    # Save Model
    saver = tf.train.Saver()
    saver.save(sess, save_path)
    print('Model Trained and Saved')

# In[32]:

plt.show()

# ### Save Parameters
# Save the `batch_size` and `save_path` parameters for inference.

# In[37]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
# Save parameters for checkpoint
helper.save_params(save_path)

# # Checkpoint

# In[38]:

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
import tensorflow as tf
import numpy as np
import helper
import problem_unittests as tests

_, (source_vocab_to_int, target_vocab_to_int), (source_int_to_vocab, target_int_to_vocab) = helper.load_preprocess()
load_path = helper.load_params()

# ## Sentence to Sequence
# To feed a sentence into the model for translation, you first need to preprocess it.  Implement the function `sentence_to_seq()` to preprocess new sentences.
#
# - Convert the sentence to lowercase
# - Convert words into ids using `vocab_to_int`
# - Convert words not in the vocabulary, to the `<UNK>` word id.

# In[39]:

def sentence_to_seq(sentence, vocab_to_int):
    """
    Convert a sentence to a sequence of ids
    :param sentence: String
    :param vocab_to_int: Dictionary to go from the words to an id
    :return: List of word ids
    """
    # TODO: Implement Function
    words_lower_case = [word.lower() for word in sentence.split()]
    # default to '<UNK>' in case word is non-existent in domain
    seq_word_id = [vocab_to_int.get(l_case_word, vocab_to_int['<UNK>']) for l_case_word in words_lower_case]

    #print(seq_word_id)
    return seq_word_id

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
tests.test_sentence_to_seq(sentence_to_seq)

# ## Translate
# This will translate `translate_sentence` from English to French.

# In[40]:

translate_sentence = 'he saw a old yellow truck .'

"""
DON'T MODIFY ANYTHING IN THIS CELL
"""
translate_sentence = sentence_to_seq(translate_sentence, source_vocab_to_int)

loaded_graph = tf.Graph()
with tf.Session(graph=loaded_graph) as sess:
    # Load saved model
    loader = tf.train.import_meta_graph(load_path + '.meta')
    loader.restore(sess, load_path)

    input_data = loaded_graph.get_tensor_by_name('input:0')
    logits = loaded_graph.get_tensor_by_name('logits:0')
    keep_prob = loaded_graph.get_tensor_by_name('keep_prob:0')

    translate_logits = sess.run(logits, {input_data: [translate_sentence], keep_prob: 1.0})[0]

print('Input')
print('  Word Ids:      {}'.format([i for i in translate_sentence]))
print('  English Words: {}'.format([source_int_to_vocab[i] for i in translate_sentence]))

print('\nPrediction')
print('  Word Ids:      {}'.format([i for i in np.argmax(translate_logits, 1)]))
print('  French Words: {}'.format([target_int_to_vocab[i] for i in np.argmax(translate_logits, 1)]))

# ## Imperfect Translation
# You might notice that some sentences translate better than others.  Since the dataset you're using only has a vocabulary of 227 English words of the thousands that you use, you're only going to see good results using these words.  Additionally, the translations in this data set were made by Google translate, so the translations themselves aren't particularly good.  (We apologize to the French speakers out there!) Thankfully, for this project, you don't need a perfect translation. However, if you want to create a better translation model, you'll need better data.
#
# You can train on the [WMT10 French-English corpus](http://www.statmt.org/wmt10/training-giga-fren.tar).  This dataset has more vocabulary and richer in topics discussed.  However, this will take you days to train, so make sure you've a GPU and the neural network is performing well on dataset we provided.  Just make sure you play with the WMT10 corpus after you've submitted this project.
# ## Submitting This Project
# When submitting this project, make sure to run all the cells before saving the notebook. Save the notebook file as "dlnd_language_translation.ipynb" and save it as a HTML file under "File" -> "Download as". Include the "helper.py" and "problem_unittests.py" files in your submission.

# In[ ]:


# -*- coding: utf-8 -*-
"""
Colour Checker Chromaticity Coordinates
=======================================

Defines *Colour Checker* chromaticity coordinates in *CIE xyY* colourspace.

Each *Colour Checker* data is in the form of a list of an :class:`OrderedDict`
class instance of 24 samples as follows::

    {'name': 'xyY', ..., 'name': 'xyY'}

The following *Colour Checkers* are available:

-   :attr:`colour.characterisation.datasets.colour_checkers.\
chromaticity_coordinates.COLORCHECKER_1976`: *Colour Checker* developed by
    *McCamy et al.* at Macbeth, a Division of Kollmorgen.
-   :attr:`colour.characterisation.datasets.colour_checkers.\
chromaticity_coordinates.COLORCHECKER_2005`: Reference data from
    *GretagMacbeth* published in 2005.
-   :attr:`colour.characterisation.datasets.colour_checkers.\
chromaticity_coordinates.BABELCOLOR_AVERAGE`: Average data derived from
    measurements of 30 *Colour Checker* charts.
-   :attr:`colour.characterisation.datasets.colour_checkers.\
chromaticity_coordinates.COLORCHECKER24_BEFORE_NOV2014`: Reference data from
    *X-Rite* published in 2015 and matching the data from *GretagMacbeth*
    published in 2005.
-   :attr:`colour.characterisation.datasets.colour_checkers.\
chromaticity_coordinates.COLORCHECKER24_AFTER_NOV2014`: Reference data from
    *X-Rite* published in 2015 and matching the *Colour Checker* edition after
    November 2014.

References
----------
-   :cite:`BabelColor2012b` : BabelColor. (2012). The ColorChecker (since
    1976!). Retrieved September 26, 2014, from
    http://www.babelcolor.com/main_level/ColorChecker.htm
-   :cite:`BabelColor2012c` : BabelColor. (2012). ColorChecker RGB and
    spectra.
    http://www.babelcolor.com/download/ColorChecker_RGB_and_spectra.xls
-   :cite:`X-Rite2016` : X-Rite. (2016). New color specifications for
    ColorChecker SG and Classic Charts. Retrieved October 29, 2018, from
    http://xritephoto.com/ph_product_overview.aspx?ID=938&Action=Support&\
SupportID=5884#
"""

from __future__ import division, unicode_literals

import numpy as np
from collections import OrderedDict, namedtuple

from colour.colorimetry import ILLUMINANTS
from colour.models import Lab_to_XYZ, XYZ_to_xyY
from colour.utilities import CaseInsensitiveMapping

__author__ = 'Colour Developers, Danny Pascale '
__copyright__ = 'Copyright (C) 2013-2020 - Colour Developers'
__copyright__ += ', '
__copyright__ += (
    'BabelColor ColorChecker data: Copyright (C) 2004-2012 Danny Pascale '
    '(www.babelcolor.com); used by permission.')
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'ColourChecker', 'COLORCHECKER_CLASSIC_SAMPLE_LABELS',
    'COLORCHECKER_1976_DATA', 'COLORCHECKER_1976_ILLUMINANT',
    'COLORCHECKER_1976', 'COLORCHECKER_2005_DATA',
    'COLORCHECKER_2005_ILLUMINANT', 'COLORCHECKER_2005',
    'BABELCOLOR_AVERAGE_DATA', 'BABELCOLOR_AVERAGE_ILLUMINANT',
    'BABELCOLOR_AVERAGE', 'COLORCHECKER24_BEFORE_NOV2014_LAB_DATA',
    'COLORCHECKER24_BEFORE_NOV2014_DATA',
    'COLORCHECKER24_BEFORE_NOV2014_ILLUMINANT',
    'COLORCHECKER24_BEFORE_NOV2014', 'COLORCHECKER24_AFTER_NOV2014_LAB_DATA',
    'COLORCHECKER24_AFTER_NOV2014_DATA',
    'COLORCHECKER24_AFTER_NOV2014_ILLUMINANT', 'COLORCHECKER24_AFTER_NOV2014',
    'COLOURCHECKERS'
]

class ColourChecker(
        namedtuple('ColourChecker', ('name', 'data', 'illuminant'))):
    """
    *Colour Checker* data.

    Parameters
    ----------
    name : unicode
        *Colour Checker* name.
    data : OrderedDict
        chromaticity coordinates in *CIE xyY* colourspace.
    illuminant : array_like
        *Colour Checker* illuminant chromaticity coordinates.
    """

COLORCHECKER_CLASSIC_SAMPLE_LABELS = (
    'dark skin',
    'light skin',
    'blue sky',
    'foliage',
    'blue flower',
    'bluish green',
    'orange',
    'purplish blue',
    'moderate red',
    'purple',
    'yellow green',
    'orange yellow',
    'blue',
    'green',
    'red',
    'yellow',
    'magenta',
    'cyan',
    'white 9.5 (.05 D)',
    'neutral 8 (.23 D)',
    'neutral 6.5 (.44 D)',
    'neutral 5 (.70 D)',
    'neutral 3.5 (1.05 D)',
    'black 2 (1.5 D)',
)
"""
*ColorChecker Classic* illuminant.

COLORCHECKER_CLASSIC_SAMPLE_LABELS : tuple
"""

COLORCHECKER_1976_DATA = OrderedDict(
    zip(COLORCHECKER_CLASSIC_SAMPLE_LABELS, [
        np.array([0.4002, 0.3504, 0.1005]),
        np.array([0.3773, 0.3446, 0.3582]),
        np.array([0.2470, 0.2514, 0.1933]),
        np.array([0.3372, 0.4220, 0.1329]),
        np.array([0.2651, 0.2400, 0.2427]),
        np.array([0.2608, 0.3430, 0.4306]),
        np.array([0.5060, 0.4070, 0.3005]),
        np.array([0.2110, 0.1750, 0.1200]),
        np.array([0.4533, 0.3058, 0.1977]),
        np.array([0.2845, 0.2020, 0.0656]),
        np.array([0.3800, 0.4887, 0.4429]),
        np.array([0.4729, 0.4375, 0.4306]),
        np.array([0.1866, 0.1285, 0.0611]),
        np.array([0.3046, 0.4782, 0.2339]),
        np.array([0.5385, 0.3129, 0.1200]),
        np.array([0.4480, 0.4703, 0.5910]),
        np.array([0.3635, 0.2325, 0.1977]),
        np.array([0.1958, 0.2519, 0.1977]),
        np.array([0.3101, 0.3163, 0.9001]),
        np.array([0.3101, 0.3163, 0.5910]),
        np.array([0.3101, 0.3163, 0.3620]),
        np.array([0.3101, 0.3163, 0.1977]),
        np.array([0.3101, 0.3163, 0.0900]),
        np.array([0.3101, 0.3163, 0.0313]),
    ]))

COLORCHECKER_1976_ILLUMINANT = (
    ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['C'])
"""
*ColorChecker 1976* illuminant.

COLORCHECKER_1976_ILLUMINANT : ndarray
"""

COLORCHECKER_1976 = ColourChecker('ColorChecker 1976', COLORCHECKER_1976_DATA,
                                  COLORCHECKER_1976_ILLUMINANT)
"""
*Colour Checker* developed by *McCamy et al.* at Macbeth, a Division of
Kollmorgen.

COLORCHECKER_1976 : ColourChecker
"""

COLORCHECKER_2005_DATA = OrderedDict(
    zip(COLORCHECKER_CLASSIC_SAMPLE_LABELS, [
        np.array([0.4316, 0.3777, 0.1008]),
        np.array([0.4197, 0.3744, 0.3495]),
        np.array([0.2760, 0.3016, 0.1836]),
        np.array([0.3703, 0.4499, 0.1325]),
        np.array([0.2999, 0.2856, 0.2304]),
        np.array([0.2848, 0.3911, 0.4178]),
        np.array([0.5295, 0.4055, 0.3118]),
        np.array([0.2305, 0.2106, 0.1126]),
        np.array([0.5012, 0.3273, 0.1938]),
        np.array([0.3319, 0.2482, 0.0637]),
        np.array([0.3984, 0.5008, 0.4446]),
        np.array([0.4957, 0.4427, 0.4357]),
        np.array([0.2018, 0.1692, 0.0575]),
        np.array([0.3253, 0.5032, 0.2318]),
        np.array([0.5686, 0.3303, 0.1257]),
        np.array([0.4697, 0.4734, 0.5981]),
        np.array([0.4159, 0.2688, 0.2009]),
        np.array([0.2131, 0.3023, 0.1930]),
        np.array([0.3469, 0.3608, 0.9131]),
        np.array([0.3440, 0.3584, 0.5894]),
        np.array([0.3432, 0.3581, 0.3632]),
        np.array([0.3446, 0.3579, 0.1915]),
        np.array([0.3401, 0.3548, 0.0883]),
        np.array([0.3406, 0.3537, 0.0311]),
    ]))

COLORCHECKER_2005_ILLUMINANT = (
    ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['ICC D50'])
"""
*ColorChecker 2005* illuminant.

COLORCHECKER_2005_ILLUMINANT : ndarray
"""

COLORCHECKER_2005 = ColourChecker('ColorChecker 2005', COLORCHECKER_2005_DATA,
                                  COLORCHECKER_2005_ILLUMINANT)
"""
Reference data from *GretagMacbeth (2005)*.

COLORCHECKER_2005 : ColourChecker
"""
BABELCOLOR_AVERAGE_DATA = OrderedDict(
    zip(COLORCHECKER_CLASSIC_SAMPLE_LABELS, [
        np.array([0.4325, 0.3788, 0.1034]),
        np.array([0.4191, 0.3748, 0.3525]),
        np.array([0.2761, 0.3004, 0.1847]),
        np.array([0.3700, 0.4501, 0.1335]),
        np.array([0.3020, 0.2877, 0.2324]),
        np.array([0.2856, 0.3910, 0.4174]),
        np.array([0.5291, 0.4075, 0.3117]),
        np.array([0.2339, 0.2155, 0.1140]),
        np.array([0.5008, 0.3293, 0.1979]),
        np.array([0.3326, 0.2556, 0.0644]),
        np.array([0.3989, 0.4998, 0.4435]),
        np.array([0.4962, 0.4428, 0.4358]),
        np.array([0.2040, 0.1696, 0.0579]),
        np.array([0.3270, 0.5033, 0.2307]),
        np.array([0.5709, 0.3298, 0.1268]),
        np.array([0.4694, 0.4732, 0.6081]),
        np.array([0.4177, 0.2704, 0.2007]),
        np.array([0.2151, 0.3037, 0.1903]),
        np.array([0.3488, 0.3628, 0.9129]),
        np.array([0.3451, 0.3596, 0.5885]),
        np.array([0.3446, 0.3590, 0.3595]),
        np.array([0.3438, 0.3589, 0.1912]),
        np.array([0.3423, 0.3576, 0.0893]),
        np.array([0.3439, 0.3565, 0.0320]),
    ]))

BABELCOLOR_AVERAGE_ILLUMINANT = (
    ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['ICC D50'])
"""
*BabelColor Average* illuminant.

BABELCOLOR_AVERAGE_ILLUMINANT : ndarray
"""

BABELCOLOR_AVERAGE = ColourChecker('BabelColor Average',
                                   BABELCOLOR_AVERAGE_DATA,
                                   BABELCOLOR_AVERAGE_ILLUMINANT)
"""
Average data derived from measurements of 30 *Colour Checker* charts.

BABELCOLOR_AVERAGE : ColourChecker
"""

COLORCHECKER24_BEFORE_NOV2014_LAB_DATA = OrderedDict(
    zip(COLORCHECKER_CLASSIC_SAMPLE_LABELS, [
        np.array([37.986, 13.555, 14.059]),
        np.array([65.711, 18.13, 17.81]),
        np.array([49.927, -4.88, -21.905]),
        np.array([43.139, -13.095, 21.905]),
        np.array([55.112, 8.844, -25.399]),
        np.array([70.719, -33.397, -0.199]),
        np.array([62.661, 36.067, 57.096]),
        np.array([40.02, 10.41, -45.964]),
        np.array([51.124, 48.239, 16.248]),
        np.array([30.325, 22.976, -21.587]),
        np.array([72.532, -23.709, 57.255]),
        np.array([71.941, 19.363, 67.857]),
        np.array([28.778, 14.179, -50.297]),
        np.array([55.261, -38.342, 31.37]),
        np.array([42.101, 53.378, 28.19]),
        np.array([81.733, 4.039, 79.819]),
        np.array([51.935, 49.986, -14.574]),
        np.array([51.038, -28.631, -28.638]),
        np.array([96.539, -0.425, 1.186]),
        np.array([81.257, -0.638, -0.335]),
        np.array([66.766, -0.734, -0.504]),
        np.array([50.867, -0.153, -0.27]),
        np.array([35.656, -0.421, -1.231]),
        np.array([20.461, -0.079, -0.973]),
    ]))
"""
*ColorChecker24 - Before November 2014* illuminant.

Notes
-----
-   *X-Rite* data is given as *CIE L\\*a\\*b\\** colourspace values under
    *CIE Illuminant D Series D50* for the
    *CIE 1931 2 Degree Standard Observer*.

COLORCHECKER24_BEFORE_NOV2014_LAB_DATA : ndarray
"""

COLORCHECKER24_BEFORE_NOV2014_DATA = OrderedDict(
    zip(
        COLORCHECKER_CLASSIC_SAMPLE_LABELS,
        XYZ_to_xyY(
            Lab_to_XYZ(
                list(COLORCHECKER24_BEFORE_NOV2014_LAB_DATA.values()),
                ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
                    'ICC D50']))))

COLORCHECKER24_BEFORE_NOV2014_ILLUMINANT = (
    ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['ICC D50'])
"""
*ColorChecker24 - Before November 2014* illuminant.

COLORCHECKER24_BEFORE_NOV2014_ILLUMINANT : ndarray
"""

COLORCHECKER24_BEFORE_NOV2014 = ColourChecker(
    'ColorChecker24 - Before November 2014',
    COLORCHECKER24_BEFORE_NOV2014_DATA,
    COLORCHECKER24_BEFORE_NOV2014_ILLUMINANT)
"""
Reference *Colour Checker* data from *X-Rite (2015)*.

Notes
-----
-   The rounded *ColorChecker24 - Before November 2014* values should match the
    *ColorChecker 2005* values. They are given for reference of the original
    *CIE L\\*a\\*b\\** colourspace values.

COLORCHECKER24_BEFORE_NOV2014 : ColourChecker
"""

COLORCHECKER24_AFTER_NOV2014_LAB_DATA = OrderedDict((
    ('dark skin', np.array([37.54, 14.37, 14.92])),
    ('light skin', np.array([64.66, 19.27, 17.5])),
    ('blue sky', np.array([49.32, -3.82, -22.54])),
    ('foliage', np.array([43.46, -12.74, 22.72])),
    ('blue flower', np.array([54.94, 9.61, -24.79])),
    ('bluish green', np.array([70.48, -32.26, -0.37])),
    ('orange', np.array([62.73, 35.83, 56.5])),
    ('purplish blue', np.array([39.43, 10.75, -45.17])),
    ('moderate red', np.array([50.57, 48.64, 16.67])),
    ('purple', np.array([30.1, 22.54, -20.87])),
    ('yellow green', np.array([71.77, -24.13, 58.19])),
    ('orange yellow', np.array([71.51, 18.24, 67.37])),
    ('blue', np.array([28.37, 15.42, -49.8])),
    ('green', np.array([54.38, -39.72, 32.27])),
    ('red', np.array([42.43, 51.05, 28.62])),
    ('yellow', np.array([81.8, 2.67, 80.41])),
    ('magenta', np.array([50.63, 51.28, -14.12])),
    ('cyan', np.array([49.57, -29.71, -28.32])),
    ('white 9.5 (.05 D)', np.array([95.19, -1.03, 2.93])),
    ('neutral 8 (.23 D)', np.array([81.29, -0.57, 0.44])),
    ('neutral 6.5 (.44 D)', np.array([66.89, -0.75, -0.06])),
    ('neutral 5 (.70 D)', np.array([50.76, -0.13, 0.14])),
    ('neutral 3.5 (1.05 D)', np.array([35.63, -0.46, -0.48])),
    ('black 2 (1.5 D)', np.array([20.64, 0.07, -0.46])),
))
"""
*ColorChecker24 - After November 2014* illuminant.

Notes
-----
-   *X-Rite* data is given as *CIE L\\*a\\*b\\** colourspace values under
    *CIE Illuminant D Series D50* for the
    *CIE 1931 2 Degree Standard Observer*.

COLORCHECKER24_AFTER_NOV2014_LAB_DATA : ndarray
"""

COLORCHECKER24_AFTER_NOV2014_DATA = OrderedDict(
    zip(
        COLORCHECKER24_AFTER_NOV2014_LAB_DATA.keys(),
        XYZ_to_xyY(
            Lab_to_XYZ(
                list(COLORCHECKER24_AFTER_NOV2014_LAB_DATA.values()),
                ILLUMINANTS['CIE 1931 2 Degree Standard Observer'][
                    'ICC D50']))))

COLORCHECKER24_AFTER_NOV2014_ILLUMINANT = (
    ILLUMINANTS['CIE 1931 2 Degree Standard Observer']['ICC D50'])
"""
*ColorChecker24 - After November 2014* illuminant.

COLORCHECKER24_AFTER_NOV2014_ILLUMINANT : ndarray
"""

COLORCHECKER24_AFTER_NOV2014 = ColourChecker(
    'ColorChecker24 - After November 2014', COLORCHECKER24_AFTER_NOV2014_DATA,
    COLORCHECKER24_AFTER_NOV2014_ILLUMINANT)
"""
Reference *Colour Checker* data from *X-Rite (2015)* and matching the
*Colour Checker* edition after November 2014.

COLORCHECKER24_AFTER_NOV2014 : ColourChecker
"""

COLOURCHECKERS = CaseInsensitiveMapping({
    'ColorChecker 1976': COLORCHECKER_1976,
    'ColorChecker 2005': COLORCHECKER_2005,
    'BabelColor Average': BABELCOLOR_AVERAGE,
    'ColorChecker24 - Before November 2014': COLORCHECKER24_BEFORE_NOV2014,
    'ColorChecker24 - After November 2014': COLORCHECKER24_AFTER_NOV2014,
})
COLOURCHECKERS.__doc__ = """
Aggregated *Colour Checker* chromaticity coordinates.

References
----------
:cite:`BabelColor2012b`, :cite:`BabelColor2012c`, :cite:`X-Rite2016`

COLOURCHECKERS : CaseInsensitiveMapping
    **{'ColorChecker 1976', 'ColorChecker 2005', 'BabelColor Average',
    'ColorChecker24 - Before November 2014',
    'ColorChecker24 - After November 2014'}**

Aliases:

-   'babel_average': 'BabelColor Average'
-   'cc2005': 'ColorChecker 2005'
-   'ccb2014': 'ColorChecker24 - Before November 2014'
-   'cca2014': 'ColorChecker24 - After November 2014'
"""
COLOURCHECKERS['babel_average'] = COLOURCHECKERS['BabelColor Average']
COLOURCHECKERS['cc2005'] = COLOURCHECKERS['ColorChecker 2005']
COLOURCHECKERS['ccb2014'] = COLOURCHECKERS[
    'ColorChecker24 - Before November 2014']
COLOURCHECKERS['cca2014'] = COLOURCHECKERS[
    'ColorChecker24 - After November 2014']


# © 2019 University of Illinois Board of Trustees.  All rights reserved
"""
Analyze variants directly from Pileup data
"""
from PileupContainerLite import PileupContainerLite
# from ReferenceCache import ReferenceCache
from PySamFastaWrapper import PySamFastaWrapper as ReferenceCache
import pickle
from functools import reduce
import numpy as np
from vcfFromContigs import createVcfRecord
import argparse
import logging
import vcf
import sys
import ast
import torch
import math
import AlleleSearcherDNN
import os
import MixtureOfExpertsTools
import MixtureOfExpertsAdvanced
import PileupDataTools
import trainDataTools
import h5py
import libCallability
from MemmapDataLite import MemmapperCompound
import pybedtools
import random

torch.set_num_threads(1);  # Torch to be run using a single thread
random.seed(13)
np.random.seed(13)

try:
    profile
except NameError:
    def profile(x):
        return x;

FEATURE_LENGTH = 150

def one_hot_encode(bases):
    """
    One-hot encode a base sequence

    :param bases: str
        Base sequence

    :return: np.array
        A list with one-hot-encoded values
    """
    indices = dict(zip('ACGT', range(4)))
    zeros = torch.zeros(len(bases), 5)
    encoded_bases = [indices[i] if (i in indices) else 4 for i in bases]
    zeros[list(range(len(bases))), encoded_bases] = 1
    return zeros

def get_reference_segment(chromosome, start, stop, ref):
    """
    Get reference slices in a given region

    :param chromosome: str
        Chromosome

    :param start: int
        Start position of variant location

    :param stop: int
        Stop position of variant location

    :param ref: PySamFastaWrapper
        Reference object

    :return: np.array
        Numpy array with one-hot encoded bases
    """
    span = FEATURE_LENGTH
    ref.chrom = chromosome
    midpoint = (start + stop) // 2
    left = midpoint - span // 2
    right = left + span
    refbases = ''.join(ref[left: right])
    return one_hot_encode(refbases)

def chromosomesInCluster(hotspots):
    """
    Determine chromosomes in a hotspots file

    :param hotspots: str
        Hotspots filename

    :return: dict
        Chromosomes and first last positions within the chromosomes in the cluster
    """
    chromosomes = dict();

    with open(hotspots, 'r') as fhandle:
        for line in fhandle:
            point = ast.literal_eval(line);
            chromosome = point['chromosome'];
            position = point['position'];
            if chromosome in chromosomes:
                values = chromosomes[chromosome];
                chromosomes[chromosome] = (min(values[0], position), max(values[1], position + 1));
            else:
                chromosomes[chromosome] = (position, position + 1);

    return chromosomes;

def number_alleles(alleles, ref_allele):
    """
    Add VCF_type numbering to alleles

    :param alleles: list
        List of alleles

    :param ref_allele: str
        Reference allele

    :return: dict
        A dictionary mapping alleles to number
    """
    numbering = dict()
    counter = 1

    for a in alleles:
        if a != ref_allele:
            numbering[a] = counter
            counter += 1

    numbering[ref_allele] = 0

    return numbering

def dump_vcfs_from_labels(entry, ref, vhandle):
    """
    Given site data, dump out a vcf file. This is to test if the site data has been
    correctly labeled or not

    :param entry: dict
        Site data dictionary

    :param ref: ReferenceCache
        Accessor for indexed references

    :param vhandle: File
        File handle into which to write VCF
    """
    logging.debug("Received item from datagen")
    site_type = 'PASS'

    if 'type' in entry:
        site_type = entry['type']
        logging.debug("Discarding entry")
        return

    alleles = entry['alleles']
    chromosome = entry['chromosome']
    ref.chrom = chromosome
    start = entry['start']
    stop = entry['stop']
    ref_allele = ''.join(ref[start: stop])
    labels = entry['labels']
    allele_numbering = number_alleles(alleles, ref_allele)
    positive_alleles = set(a for a, l in zip(alleles, labels) if l > 0)
    genotypes = [allele_numbering[a] for a in positive_alleles]
    alt_alleles = [a for a in alleles if a != ref_allele]

    if len(genotypes) == 0:
        genotypes = [0, 0]
        site_type = 'UNLABELED'

    if len(genotypes) == 1:
        genotypes += genotypes

    if len(alt_alleles) == 0:
        alt_alleles = [ref_allele]
        site_type = 'NO_ALT'

    logging.debug("Creating vcf record")

    record = createVcfRecord(
        chromosome,
        start,
        ref,
        [0],
        [ref_allele],
        [alt_alleles],
        [genotypes],
        string="LabeledFromGroundTruth",
        qual=50,
        qualifier=site_type
    )

    logging.debug("Completed ceating vcf record")

    for r in record:
        vhandle.write(r + '\n')

def intersect(vcf, bed, hotspots, outputPrefix):
    """
    Given a hotspots file, determine minimally spanning bed/vcf files

    :param vcf: str
        Vcf file

    :param bed: str
        Bed file

    :param hotspots: str
        hotspots file

    :param outputPrefix: str
        Prefix of output file

    :return: tuple
        (new vcf, new bed)
    """
    bedName = outputPrefix + ".bed";
    vcfName = outputPrefix + ".vcf";
    chromosomes = chromosomesInCluster(hotspots);
    bedString = "";

    for i, (key, value) in enumerate(chromosomes.items()):
        chromosome = key;
        start = value[0] - 10;
        stop = value[1] + 10;
        bedString += "\n" if i > 0 else "";
        bedString += "\t".join([chromosome, str(start), str(stop)]);

    subset = pybedtools.BedTool(bedString, from_string=True);
    vcfBed = pybedtools.BedTool(vcf);
    hcBed = pybedtools.BedTool(bed);
    vcfBed.intersect(subset, header=True).saveas(vcfName);
    hcBed.intersect(subset).saveas(bedName);

    return (vcfName, bedName);

class NoReadsFoundError(Exception):
    """
    Custom exception thrown by PileupAnalyzer indicating that
    there are no reads at the site
    """
    def __init__(self, value):
        self.value = value;

    def __str__(self):
        return repr(self.value);

def copy(inhandle, outhandle, location, exclude_keys):
    outgroup = outhandle.create_group(location);
    ingroup = inhandle[location];

    for allele in ingroup.keys():
        if allele == 'siteLabel':
            dset = outgroup.create_dataset('siteLabel', dtype=np.int32, shape=(1, ));
            dset[:] = ingroup['siteLabel'][:];
            continue;

        alleleInGroup = ingroup[allele];
        alleleOutGroup = outgroup.create_group(allele);

        for attribute in alleleInGroup.keys():
            if attribute in exclude_keys:
                continue
            dset = alleleOutGroup.create_dataset(
                attribute, dtype=alleleInGroup[attribute].dtype, shape=alleleInGroup[attribute].shape
            );
            dset[:] = alleleInGroup[attribute][:];

def postProcessHdf5(
    hdf5name,
    outputPrefix,
    hybrid,
):
    """
    For generating training data, convert the pure hdf5 file into .index format

    :param hdf5name: str
        Name of the hdf5 file

    :param outputPrefix: str
        Prefix of the output file

    :param hybrid: bool
        Whether we are in hybrid mode or not
    """
    # Anand: Added reference segment on September 14 2020
    memmapperKeys = ['feature', 'feature2', 'label', 'segment'];
    memmapperConcatenations = [False, False, True, False];

    if not hybrid:
        # Anand: Added reference segment on September 14 2020
        memmapperKeys = ['feature', 'label', 'segment'];
        memmapperConcatenations = [False, True, False];

    # Anand: Added reference segment on September 14 2020
    dtypes = {'feature': 'uint8', 'feature2': 'uint8', 'segment': 'uint8', 'label': 'float32'};

    logging.info("Converting to integrated format");

    MemmapperCompound(
        hdf5name,
        prefix=outputPrefix,
        keys=memmapperKeys,
        concatenations=memmapperConcatenations,
        hybrid=hybrid,
        mode='integrated',
        dtypes=dtypes,
        memtype="memmap",
    );

def addToDict(
    fhandle,
    chromosome,
    start,
    stop,
    alleles,
    tensors,
    labels,
    siteLabel,
    scores,
    supportingReads,
    supportingReadsStrict,
    tensors2,
    supportingReads2,
    supportingReadsStrict2,
    segment,
):
    """
    Add contents to an hdf5 file

    :param fhandle: dict
        Dictionary data into which to add items

    :param chromosome: str
        Chromosome

    :param start: int
        Start position in chromosome reference

    :param stop: int
        Stop position in chromosome reference

    :param alleles: list
        List of alleles at site

    :param tensors: list
        Feature tensors (np.ndarray)

    :param labels: list
        List of labels in the same order as tensors

    :param siteLabel: int
        Whether site is heterozygous or homozygous

    :param scores: list
        Scoring results from network

    :param supportingReads: int
        Number of supporting reads for each allele

    :param supportingReadsStrict: int
        Number of reads containing the allele in its entirety

    :param tensors2: list
        Second set of feature tensors for hybrid variant calling

    :param supportingReads2: list
        Second set of number of supporting reads for hybrid calling

    :param supportingReadsStrict2:
        Second set of strict supporting reads for hybrid calling

    :param segment: np.array
        Reference segment
    """
    logging.debug("Writing chr, start, stop = %s, %d, %d to file" % (chromosome, start, stop));
    groupName = '_'.join([chromosome, str(start), str(stop)]);  # We will maintain a flat hierarchy

    if groupName in fhandle:
        del fhandle[groupName];
    else:
        fhandle[groupName] = dict()
        mainGroup = fhandle[groupName]

        for i, (allele, tensor, label) in enumerate(zip(alleles, tensors, labels)):
            mainGroup[allele] = dict()
            alleleGroup = mainGroup[allele];
            alleleGroup['label'] = np.zeros(shape=(1, ), dtype=np.float32)
            alleleGroup['feature'] = np.zeros(shape=tensor.shape, dtype=tensor.dtype)
            alleleGroup['label'][:] = label;
            alleleGroup['feature'][:] = tensor;
            alleleGroup['supportingReads'] = [supportingReads[i]];
            alleleGroup['supportingReadsStrict'] = [supportingReadsStrict[i]];

            # Adding reference segment: adding it to each allele (September 14 2020)
            alleleGroup['segment'] = np.zeros(shape=segment.shape, dtype=np.uint8)
            alleleGroup['segment'][:] = segment[:]

            if (len(tensors2) > 0):
                tensor2 = tensors2[i];
                alleleGroup['feature2'] = np.zeros(shape=tensor2.shape, dtype=tensor2.dtype)
                alleleGroup['feature2'][:] = tensor2;
                alleleGroup['supportingReads2'] = [supportingReads2[i]];
                alleleGroup['supportingReadsStrict2'] = [supportingReadsStrict2[i]];

        mainGroup['siteLabel'] = siteLabel

def addToHDF5File(
    fhandle,
    chromosome,
    start,
    stop,
    alleles,
    tensors,
    labels,
    siteLabel,
    scores,
    supportingReads,
    supportingReadsStrict,
    tensors2,
    supportingReads2,
    supportingReadsStrict2,
    segment,
):
    """
    Add contents to an hdf5 file

    :param fhandle: h5py.File
        File object into which a group should be added

    :param chromosome: str
        Chromosome

    :param start: int
        Start position in chromosome reference

    :param stop: int
        Stop position in chromosome reference

    :param alleles: list
        List of alleles at site

    :param tensors: list
        Feature tensors (np.ndarray)

    :param labels: list
        List of labels in the same order as tensors

    :param siteLabel: int
        Whether site is heterozygous or homozygous

    :param scores: list
        Scoring results from network

    :param supportingReads: int
        Number of supporting reads for each allele

    :param supportingReadsStrict: int
        Number of reads containing the allele in its entirety

    :param tensors2: list
        Second set of feature tensors for hybrid variant calling

    :param supportingReads2: list
        Second set of number of supporting reads for hybrid calling

    :param supportingReadsStrict2:
        Second set of strict supporting reads for hybrid calling

    :param segment: np.array
        Reference segment
    """
    logging.debug("Writing chr, start, stop = %s, %d, %d to file" % (chromosome, start, stop));
    groupName = '_'.join([chromosome, str(start), str(stop)]);  # We will maintain a flat hierarchy

    try:
        # If group already exists, delete it: it might be erroneous
        groupData = fhandle[groupName];
        del fhandle[groupName];
    except Exception:
        # If group doesn't exist, add data
        mainGroup = fhandle.create_group(groupName);

        for i, (allele, tensor, label) in enumerate(zip(alleles, tensors, labels)):
            alleleGroup = mainGroup.create_group(allele);
            alleleGroup.create_dataset('label', shape=(1,), dtype='float32');
            alleleGroup.create_dataset('feature', shape=tensor.shape, dtype=tensor.dtype);
            alleleGroup.create_dataset('supportingReads', shape=(1,), dtype='int32');
            alleleGroup.create_dataset('supportingReadsStrict', shape=(1,), dtype='int32');
            alleleGroup['label'][:] = label;
            alleleGroup['feature'][:] = tensor;
            alleleGroup['supportingReads'][:] = supportingReads[i];
            alleleGroup['supportingReadsStrict'][:] = supportingReadsStrict[i];

            # Adding reference segment: adding it to each allele (September 14 2020)
            alleleGroup.create_dataset('segment', shape=segment.shape, dtype='uint8')
            alleleGroup['segment'][:] = segment[:]

            if (len(tensors2) > 0):
                tensor2 = tensors2[i];
                alleleGroup.create_dataset('feature2', shape=tensor2.shape, dtype='float32');
                alleleGroup.create_dataset('supportingReads2', shape=(1,), dtype='int32');
                alleleGroup.create_dataset('supportingReadsStrict2', shape=(1,), dtype='int32');
                alleleGroup['feature2'][:] = tensor2;
                alleleGroup['supportingReads2'][:] = supportingReads2[i];
                alleleGroup['supportingReadsStrict2'][:] = supportingReadsStrict2[i];

        mainGroup.create_dataset('siteLabel', shape=(1,), dtype=int);
        mainGroup['siteLabel'][:] = siteLabel;

def dumpTrainingData(datagen, filename, hybrid=False, dump_vcf=False, ref=None, keep=False):
    """
    Dump training data to disk

    :param datagen: trainDataTools.data
        Iterator from a generator

    :param filename: str
        Prefix of the output file

    :param hybrid: bool
        Whether we are operating in hybrid mode or not

    :param dump_vcf: bool
        Dump vcf file from labeled ground-truth to test labeling accuracy

    :param ref: PySamFastaWrapper
        Reference reader

    :param keep: bool
        Keep HDF5 data in original format
    """
    missedCalls = [];
    tooLong = [];
    hdf5 = filename + ".hdf5";
    vcf = filename + ".labeled.vcf"

    if dump_vcf:
        vhandle = open(vcf, 'w')

    if keep:
        fhandle = h5py.File(filename + ".hdf5", 'w')
    else:
        fhandle = dict()

    for i, entry in enumerate(datagen):
        if 'type' in entry:
            if entry['type'] == "TOO_LONG":
                tooLong.append(entry);
            else:
                missedCalls.append(entry);
        else:
            # Update Sept 14 2020 - add reference segment here
            entry['segment'] = get_reference_segment(
                entry['chromosome'],
                entry['start'],
                entry['stop'],
                ref,
            )

            entry['fhandle'] = fhandle;

            if keep:
                addToHDF5File(**entry);
            else:
                addToDict(**entry)

        if (i + 1) % 100 == 0:
            logging.info("Completed %d locations" % (i + 1));

        dump_vcfs_from_labels(entry, ref, vhandle)

    if not keep:
        postProcessHdf5(fhandle, filename, hybrid=hybrid)
    else:
        fhandle.close()

    if dump_vcf:
        vhandle.close()

def scoreSite(siteDict, network, cache, hybrid=False):
    """
    Scores the alleles at a site based on data

    :param siteDict: dict
        Tensors and other information at a site

    :param network: torch.nn.Module
        Neural network to compute likelihoods

    :param cache: PySamFastaWrapper
        A Reference object

    :param hybrid: bool
        Whether we are running the tool in hybrid mode
    """
    featureDict = dict();

    for i, (allele, t1) in enumerate(zip(siteDict['alleles'], siteDict['tensors'])):
        if hybrid:
            tensorPacket = (torch.Tensor(t1), torch.Tensor(siteDict['tensors2'][i]));
        else:
            tensorPacket = (torch.Tensor(t1), None);

        featureDict[allele] = tensorPacket;

    # Find reference (September 14 2020 - update for using reference segment)
    ref_segment = get_reference_segment(
        siteDict['chromosome'],
        siteDict['start'],
        siteDict['stop'],
        cache,
    )
    ref_segment = torch.unsqueeze(ref_segment, dim=0)

    with torch.no_grad():
        logLikelihoods = network(featureDict, ref_segment);

    return logLikelihoods;

def vcfRecords(siteDict, network, ref, hybrid=False):
    """
    Creates vcf records for a site

    :param siteDict: dict
        Tensors and other information at a site

    :param network: torch.nn.Module
        Neural network to compute likelihoods

    :param ref: ReferenceCache
        Reference cache for obtaining reference allele etc

    :param hybrid: bool
        Whether we are running the tool in hybrid mode
    """
    if ref.chrom != siteDict['chromosome']:
        ref.chrom = siteDict['chromosome'];

    supportMsg = dict(
        zip(
            siteDict['alleles'],
            zip(siteDict['supportingReadsStrict'], siteDict['supportingReadsStrict2'])
        )
    );

    logging.debug(
        "Supports at site = %s" % str(supportMsg)
    );

    returns = scoreSite(siteDict, network, ref, hybrid);

    if type(returns) is dict:
        logLikelihoods = returns;
        features = None;
    else:
        logLikelihoods = returns[0];
        features = returns[1:];

    logging.debug("Likelihoods at site = %s" % str(logLikelihoods));

    refAllele = ''.join(ref[siteDict['start']: siteDict['stop']]);
    topAlleleCombination = sorted([(v, k) for k, v in logLikelihoods.items()], reverse=True)[0];
    logLikelihood, topAlleles = topAlleleCombination;
    logLikelihood = min(float(logLikelihood), 1 - 1e-8);  # Quality score restricted to value 80
    quality = -10 * math.log10(1 - logLikelihood);
    altAlleles = list(set(topAlleles).difference({refAllele}));
    allelesAtSite = siteDict['alleles'];

    if len(altAlleles) == 0:
        genotypes = [0, 0];
        allAlleles = allelesAtSite;
        altAlleles = list(set(allAlleles).difference({refAllele}));
        if len(altAlleles) == 0:
            return None;
    else:
        genotypes = [];

        for allele in topAlleles:
            if allele == refAllele:
                genotypes.append(0);
            else:
                try:
                    altIndex = altAlleles.index(allele);
                except ValueError:
                    # Rethrow with message (should probably put in a finally block)
                    raise ValueError("Cannot find allele %s in altAllele list %s" % (allele, str(altAlleles)));

                genotypes.append(altIndex + 1);

    record = createVcfRecord(
        siteDict['chromosome'],
        siteDict['start'],
        # ref.database,
        ref,
        [0],
        [refAllele],
        [altAlleles],
        [genotypes],
        string="MixtureOfExpertPrediction",
        qual=quality
    )[0];

    if features is None:
        return record;
    else:
        expert0, expert1, expert2, meta = features;
        features = dict({
            'chromosome': siteDict['chromosome'],
            'position': siteDict['start'],
            'length': len(refAllele),
            'meta': meta.cpu().data.numpy(),
            'expertPredictions': (expert0, expert1, expert2),
        });
        return record, features;

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Call variants for Mixture of Experts model");

    parser.add_argument(
        "--bam",
        help="Comma-separated list of BAM files",
        required=True,
    );

    parser.add_argument(
        "--activity",
        help="Active regions file",
        required=True,
    );

    parser.add_argument(
        "--ref",
        help="Reference cache location",
        required=True,
    );

    parser.add_argument(
        "--network",
        help="AlleleSearcherDNN trained model path",
        required=False,
    );

    parser.add_argument(
        "--outputPrefix",
        help="Prefix of output file",
        required=True,
    );

    parser.add_argument(
        "--debug",
        help="Enable debug mode",
        action="store_true",
        default=False,
    );

    parser.add_argument(
        "--provideFeatures",
        help="Provide features in addition to vcf records",
        default=False,
        action="store_true",
    );

    parser.add_argument(
        "--hotspotMode",
        help="Mode for hotspot creation",
        choices=["BAM1", "BOTH", "INCLUSIVE"],
        default="BOTH",
    );

    parser.add_argument(
        "--chrPrefixes",
        help="Comma-separated list of prefix for BAM files",
        default=None,
    );

    parser.add_argument(
        "--featureLength",
        help="Length of feature maps",
        default=100,
        type=int,
    );

    parser.add_argument(
        "--truth",
        help="Ground truth VCF if training data generation is desired",
        default=None,
    );

    parser.add_argument(
        "--highconf",
        help="High-confidence bed file if generating training data is desired",
        default=None,
    );

    parser.add_argument(
        "--intersectRegions",
        help="Intersect ground-truth/high-confidence regions with hotspot range before running",
        default=False,
        action="store_true",
    );

    parser.add_argument(
        "--simple",
        help="Use simple features",
        default=False,
        action="store_true",
    );

    parser.add_argument(
        "--reuseSearchers",
        help="Re-use searchers between hotspots and feature generation",
        default=False,
        action="store_true",
    );

    parser.add_argument(
        "--noAlleleLevelFilter",
        help="Do not use allele-level filters for reads",
        default=False,
        action="store_true",
    );

    parser.add_argument(
        "--clr",
        action="store_true",
        help="Second read type is CLR instead of CCS",
        default=False,
    );

    parser.add_argument(
        "--hybrid_hotspot",
        help="Use hybrid hotspot detection method",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "--pacbio",
        help="For single BAM mode indicates whether a read is pacbio or not",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "--test_labeling",
        help="This mode uses ground-truth variants etc, and dumps out a vcf file to test whether labeling is correct",
        action="store_true",
        default=False,
    )

    parser.add_argument(
        "--only_contained",
        help="For training data dump, only use locations that are fully contained in confident regions",
        default=False,
        action="store_true"
    )

    parser.add_argument(
        "--hybrid_eval",
        help="Evaluate ground-truth in hybrid manner",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "--q_threshold",
        help="Quality score threshold",
        default=10,
        type=int,
    )

    parser.add_argument(
        "--mapq_threshold",
        help="Mapping quality threshold",
        default=10,
        type=int,
    )

    parser.add_argument(
        "--keep_hdf5",
        help="Retain data as hdf5, do not convert to memmap",
        default=False,
        action="store_true",
    )

    parser.add_argument(
        "--reconcilement_size",
        help="Size of a hotspot region to enable reconcilement of pacbio/illumina representations",
        default=10,
        type=int,
    )

    args = parser.parse_args();

    FEATURE_LENGTH = args.featureLength

    if args.only_contained:
        trainDataTools.STRICT_INTERSECTION = True

    if args.hybrid_eval:
        trainDataTools.HYBRID_TRUTH_EVAL = True

    libCallability.initLogging(args.debug);
    logging.basicConfig(
        level=logging.INFO if not args.debug else logging.DEBUG, format='%(asctime)-15s %(message)s'
    );

    if args.intersectRegions and (args.highconf is not None):
        logging.info("Shrinking ground-truth to fit the given active regions");
        newVcf, newBed = intersect(args.truth, args.highconf, args.activity, args.outputPrefix);
        args.truth = newVcf;
        args.highconf = newBed;
        logging.info("Obtained new vcf and bed files %s, %s" % (args.truth, args.highconf));

    # Create read factories and searcher factories
    bamfiles = args.bam.split(",");
    bamReaders = [];
    chrPrefixes = ["" for _ in bamfiles];

    if args.chrPrefixes is not None:
        chrPrefixes = args.chrPrefixes.split(",");
        assert(len(chrPrefixes) == len(bamfiles)), \
            "When provided, the length of chr prefixes should be the same as that of the bam files";

    readSamplers = [];

    for i, (bam, chrPrefix) in enumerate(zip(bamfiles, chrPrefixes)):
        reader = PileupDataTools.ReadSampler(
            bamfile=bam,
            readRate=(PileupDataTools.READ_RATE_ILLUMINA if i == 0 else PileupDataTools.READ_RATE_PACBIO),
            chrPrefix=chrPrefix,
            pacbio=(i > 0) or (len(bamfiles) == 1 and args.pacbio),
        );
        readSamplers.append(reader);

    searcherFactory = PileupDataTools.SearcherFactory(
        ref=args.ref,
        featureLength=args.featureLength,
        pacbio=args.pacbio,
        useInternalLeftAlignment=False,
        noAlleleLevelFilter=args.noAlleleLevelFilter,
        clr=False,      # TBD: Currently doesn't support CLR reads
        hybrid_hotspot=args.hybrid_hotspot,
        min_mapq=args.mapq_threshold,
        q_threshold=args.q_threshold,
        reassembly_size=args.reconcilement_size,
    );

    logging.info("Getting callable sites");

    # Obtain the list of hotspots to be called
    hotspots, searcherCollection = PileupDataTools.candidateReader(
        readSamplers=readSamplers,
        searcherFactory=searcherFactory,
        activity=args.activity,
        hotspotMode=args.hotspotMode,
        provideSearchers=args.reuseSearchers,
    );

    logging.info("Obtained %d hotspots" % len(hotspots));

    # Generate data, and either dump to disk, or call variants
    datagen = trainDataTools.data(
        hotspots,
        readSamplers,
        searcherFactory,
        args.ref,
        vcf=args.truth,
        bed=args.highconf,
        hotspotMethod=args.hotspotMode,
        searcherCollection=searcherCollection,
    );

    featureList = None;

    if args.highconf is not None:
        # if not args.test_labeling:
        #     dumpTrainingData(datagen, args.outputPrefix, hybrid=(len(bamfiles) > 1));
        # else:
        #     cache = ReferenceCache(database=args.ref);
        #     dump_vcfs_from_labels(datagen, args.outputPrefix + ".labeled", cache)
        ref = ReferenceCache(database=args.ref)
        dumpTrainingData(datagen, args.outputPrefix, hybrid=(len(bamfiles) > 1), dump_vcf=True, ref=ref, keep=args.keep_hdf5);
    else:
        trainDataTools.MAX_ITEMS_PER_GROUP = 1024
        assert(args.network is not None), "Provide DNN for variant calling";
        cache = ReferenceCache(database=args.ref);
        fhandle = open(args.outputPrefix + ".vcf", 'w');
        network = torch.load(args.network, map_location='cpu');
        network.eval();

        if args.provideFeatures:
            network.providePredictions = True;
            featureList = [];
        else:
            featureList = None;

        for i, siteDict in enumerate(datagen):
            if siteDict is None:
                continue;

            if len(siteDict['alleles']) == 0:
                continue;

            record = vcfRecords(siteDict, network, cache, hybrid=(len(bamfiles) > 1));

            if record is not None:
                if type(record) is tuple:
                    record, featureDict = record;
                    fhandle.write(str(record) + '\n');
                    featureList.append(featureDict);
                else:
                    fhandle.write(str(record) + "\n");

            if (i + 1) % 100 == 0:
                logging.info("Completed %d sites" % (i + 1));

        fhandle.close();

    if featureList is not None:
        with open(args.outputPrefix + ".features", "wb") as whandle:
            pickle.dump(featureList, whandle);

    logging.info("Completed running the script");