from __future__ import print_function
from __future__ import division
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import torch.optim as optim
import cPickle as pickle

class SaveVal():
    def __init__(self):
        self.testLoss = []
        self.trainLoss = []
        self.trainAcc = []
        self.testAcc = []

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv5 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv6 = nn.Conv2d(128, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)
        # Part 1
        self.batchNorm = nn.BatchNorm1d(512)
        # Part 2
        self.fc_between = nn.Linear(512, 512)  # part 2

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.pool(x)
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = x.view(-1, self.num_flat_features(x))
        # Part 1: Add batchnorm after fc1 and before ReLU
        x = F.relu(self.batchNorm(self.fc1(x)))
        # x = F.relu(self.fc1(x))

        # Part 2: add fc layer
        x = F.relu(self.fc_between(x))

        x = self.fc2(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features

def eval_net(dataloader):
    correct = 0
    total = 0
    total_loss = 0
    net.eval() # Why would I do this? -> because the network eval and train mode are different
    criterion = nn.CrossEntropyLoss(size_average=False)
    for data in dataloader:
        images, labels = data
        images, labels = Variable(images), Variable(labels)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels.data).sum()
        loss = criterion(outputs, labels)
        total_loss += loss.data[0]
    net.train() # Why would I do this? -> because the network eval and train mode are different
    return total_loss / total, correct / total

if __name__ == "__main__":
    # FOR SAVING
    save = SaveVal()

    BATCH_SIZE = 32 #mini_batch size
    MAX_EPOCH = 15 #maximum epoch to train

    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) #torchvision.transforms.Normalize(mean, std)

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,
                                              shuffle=True, num_workers=2)

    testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,
                                             shuffle=False, num_workers=2)

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    print('Building model...')
    net = Net()

    # **********************************************************************************************************************
    #                                       Loading Pretrained Model upto a certain layer
    # **********************************************************************************************************************
    # # Loading pickle
    # pretrained_dict = torch.load("part1_BN.pth")
    # # Current model's dict
    # model_dict = net.state_dict()
    # # Filtering keys
    # pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    # # 2. overwrite entries in the existing state dict
    # model_dict.update(pretrained_dict)
    # # 3. load the new state dict
    # net.load_state_dict(model_dict)
    # ******************************************************************************************************************

    net.train() # Training mode

    criterion = nn.CrossEntropyLoss()
    # optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
    optimizer = optim.Adam(net.parameters()) # part 3

    print('Start training...')
    for epoch in range(MAX_EPOCH):  # loop over the dataset multiple times

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # get the inputs
            inputs, labels = data

            # wrap them in Variable
            inputs, labels = Variable(inputs), Variable(labels)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            # print statistics
            running_loss += loss.data[0]
            if i % 500 == 499:    # print every 2000 mini-batches
                print('    Step: %5d avg_batch_loss: %.5f' %
                      (i + 1, running_loss / 500))
                running_loss = 0.0
        print('    Finish training this EPOCH, start evaluating...')
        train_loss, train_acc = eval_net(trainloader)
        test_loss, test_acc = eval_net(testloader)
        print('EPOCH: %d train_loss: %.5f train_acc: %.5f test_loss: %.5f test_acc %.5f' %
              (epoch+1, train_loss, train_acc, test_loss, test_acc))
        # saving losses and accuracies
        save.trainLoss.append(train_loss)
        save.trainAcc.append(train_acc)
        save.testLoss.append(test_loss)
        save.testAcc.append(test_acc)

    pickle.dump(save, open("part4_save_test_train_loss_acc_BN_addedFC_Adam_ep15_addedConv.p", "wb"))

    print('Finished Training')
    print('Saving model...')
    torch.save(net.state_dict(), 'part4_BN_addedFC_Adam_ep15_addedConv.pth')


from django.shortcuts import render, HttpResponse, redirect, get_object_or_404
from django.http import JsonResponse
from django.urls import reverse_lazy
from rest_framework.renderers import JSONRenderer
from django.db import models
from django.contrib.contenttypes.models import ContentType

from data_importation.views import delete
from data_reference.forms import *
from data_importation.forms import DeleteForm

def get_fields_for_model(model):
    model = ContentType.objects.get(model__icontains=model)
    fields = [[v.name, str([t.name for t in v.related_model._meta.fields]) if v.is_relation else False] for i, v in
              enumerate(model.model_class()._meta.fields)]
    return fields

def get_model_fields(request, model, app_label=None):
    fields = get_fields_for_model(model)
    # serial = JSONRenderer().render(fields)
    return JsonResponse(fields, safe=False, status=200)

def get_class_lookups(request, model, field_name):
    model = ContentType.objects.get(model__contains=model)
    fields = model.model_class()._meta.fields
    lookups = []
    for field in fields:
        temp = field.get_lookups().keys() if field.attname == field_name else None
        if temp is not None:
            lookups = list(temp)

    return JsonResponse(lookups, safe=False, status=200)

def get_all_models(request, app_label):
    tables = ContentType.objects.filter(app_label__icontains=app_label)
    tables = [i.model for i in tables]
    return JsonResponse(tables, safe=False, status=200)

def get_model_filter_kwargs(request, fields):
    pass

def create_model_table_view(request):
    form = ModelTableForm()

    if request.method == "POST":
        form = ModelTableForm(request.POST)

        if form.is_valid():
            form.save()

            if form.cleaned_data['next_action'] == "1":
                return redirect(reverse_lazy("data-reference:create-model-table"))
            elif form.cleaned_data["next_action"] == "2":
                return HttpResponse("Ajouter des colonnes à ce modèle")
            return HttpResponse("Form Validated!")
        else:
            return render(request, 'data_reference/create-model-table-view.html', locals())

    return render(request, 'data_reference/create-model-table-view.html', locals())

def all_model_tables(request):
    all_tables = ModelTable.objects.all()
    return render(request, 'data_reference/all_models_tables.html', locals())

def edit_model_table(request, pk):
    model_table = get_object_or_404(ModelTable, pk=pk)
    form = ModelTableForm(instance=model_table)
    if request.method == "POST":
        form = ModelTableForm(request.POST)
        if form.is_valid():
            # form.save()
            ModelTable.objects.filter(pk=pk).update(**{key:value for key, value in form.cleaned_data.items()
                                                       if key != "next_action"})

            if form.cleaned_data['next_action'] == "1":
                return redirect(reverse_lazy("data-reference:create-model-table"))
            elif form.cleaned_data["next_action"] == "2":
                return HttpResponse("Ajouter des colonnes à ce modèle")
            return redirect(reverse_lazy('data-reference:all-models-tables'))
        else:
            return render(request, 'data_reference/create-model-table-view.html', locals())

    return render(request, 'data_reference/create-model-table-view.html', locals())

def delete_model_table(request, pk):
    model_table = get_object_or_404(ModelTable, pk=pk)
    if request.method == 'POST':
        form = DeleteForm(request.POST)
        delete(request, model_table, form)
    return redirect(reverse_lazy("data-reference:all-models-tables"))

def detail_model_table(request, pk):
    model_table = get_object_or_404(ModelTable, pk=pk)
    return render(request, 'data_reference/detail-model-table.html', locals())

def create_column_for_model(request, pk):
    model_table = get_object_or_404(ModelTable, pk=pk)
    fields = get_fields_for_model(model_table.model_name)
    form = ModelColumnForm()

    if request.method == 'POST':
        form = ModelColumnForm(request.POST)

        if form.is_valid():
            instance = form.save(commit=False)
            instance.model_table = model_table
            instance.save()

            if form.cleaned_data['next_action'] == "1":
                return redirect(reverse_lazy("data-reference:create-column-for-model", kwargs={'pk': model_table.pk}))
            elif form.cleaned_data["next_action"] == "2":
                return HttpResponse("Ajouter des colonnes à ce modèle")

            return redirect(reverse_lazy("data-reference:detail-model-table", kwargs={'pk': model_table.pk}))
        else:
            return render(request, 'data_reference/create-column-for-table.html', locals())
    return render(request, 'data_reference/create-column-for-table.html', locals())

def delete_column_for_model(request, pk):
    column = get_object_or_404(ModelColumn, pk=pk)
    if request.method == 'POST':
        form = DeleteForm(request.POST)
        delete(request, column, form)
    return redirect(reverse_lazy("data-reference:detail-model-table", kwargs={'pk': column.model_table.pk}))

def edit_column_for_model(request, pk):
    pass

def view_table_by_model(request, pk):
    model_table = get_object_or_404(ModelTable, pk=pk)
    model = ContentType.objects.get(model=model_table.model_name)
    headers = [i.name for i in model_table.table_to_fields.all()]
    data = [[getattr(obj, i) for i in headers if hasattr(obj, i)] for obj in model.model_class().objects.all()]
    try:
        headers = headers[:len(data[0]) if len(headers) > len(data[0]) else len(headers)]
    except IndexError as e:
        pass

    return render(request, 'data_reference/view-table-by-model.html', locals())


import requests
import json
from airflow.models import Variable
import logging
from googletrans import Translator
from vocab_utils.translate import translate_vocab

log = logging.getLogger(__name__)

def get_definitions(vocabs: list, vocab_origins: list, target_lang: str):
    """
    get_definitions()
        Using LinguaAPI, the definitions, examples, synonyms and contexts are gathered.
        Then they are stored into a dictionary format.

    """
    vocab_dic = {}
    for ind, vocab in enumerate(vocabs):

        # DEFINE vocab_info_dic
        # try: Some vocabularies do not have definitions (ex: fugazi)
        url = "https://lingua-robot.p.rapidapi.com/language/v1/entries/en/" + \
            vocab.lower().strip(' ')

        headers = {
            "X-RapidAPI-Key": Variable.get("lingua_credentials"),
            "X-RapidAPI-Host": "lingua-robot.p.rapidapi.com"
        }

        # Request Data
        response = requests.request("GET", url, headers=headers)
        data = json.loads(response.text)

        try:
            vocab_dat = data['entries'][0]['lexemes']
        except IndexError:
            vocab_dat = None
            definitions = None
            synonyms = None
            examples = None
            audio_url = None

        def extract_audio_url(json_data):

            try:
                json_data = data['entries']
                for entry in json_data:
                    pronunciations = entry.get('pronunciations', [])
                    for pronunciation in pronunciations:
                        audio = pronunciation.get('audio')
                        if audio:
                            return audio['url']
            except IndexError:
                return None
            return None

        # GET DEFINITIONS
        # try: If the definition is not in Lingua Dictionary, output None
        try:
            definitions = [vocab_dat[j]['senses'][i]['definition']
                            for j in range(len(vocab_dat)) for i in range(len(vocab_dat[j]['senses']))]
            definitions = definitions[:5]
        except:
            pass

        # GET AUDIO URLS
        audio_url = extract_audio_url(data)

        # GET SYNONYMS
        # try: If synonyms are not in Lingua Dictionary, output None
        try:
            synonyms = [vocab_dat[j]['synonymSets'][i]['synonyms']
                        for j in range(len(vocab_dat)) for i in range(len(vocab_dat[j]['synonymSets']))]
        except:
            synonyms = None

        # GET EXAMPLES
        if vocab_origins[ind] != None:
            context = vocab_origins[ind]
        else:
            context = None

        try:
            examples += [vocab_dat[j]['senses'][i]['usageExamples']
                        for j in range(len(vocab_dat)) for i in range(len(vocab_dat[j]['senses']))
                        if 'usageExamples' in vocab_dat[j]['senses'][i].keys()]
        except:
            examples = None

        # Collect vocab details
        vocab_info_dic = {vocab: {'definitions': definitions,
                                  'examples': examples,
                                  'synonyms': synonyms,
                                  'context': context,
                                  'audio_url': audio_url}}

        if target_lang != 'en':
            translator = Translator()
            vocab_info_dic = translate_vocab(translator, vocab_info_dic, 'en', target_lang)

        vocab_dic.setdefault(vocab, []).append(vocab_info_dic[vocab])

    return vocab_dic


from django.contrib import admin

# Register your models here.
from course.models import course_category,course,lesson,course_requirements,access_to,course_comment,course_comment_reply,enrollment
from django.utils.html import format_html

admin.site.register(course_category)

class display_course(admin.ModelAdmin):
    list_display = ('course_name','photo_pre')
    readonly_fields = ('photo_pre',)
    list_per_page = 10
    search_fields = ('course_name','course_tittle')
    def photo_pre(self,obj):
        return format_html(f'')
admin.site.register(course,display_course)

class display_lesson(admin.ModelAdmin):
    list_display = ('lesson_name','photo_pre')
    readonly_fields = ('photo_pre',)
    list_per_page = 10
    search_fields = ('lesson_name','course_name')
    list_filter = ('course_name',)
    def photo_pre(self,obj):
        return format_html(f'')
admin.site.register(lesson,display_lesson)

class display_req(admin.ModelAdmin):
    list_display = ('course_name','requirement')
    list_per_page = 10
    search_fields = ('course_name','requirement')
    list_filter = ('course_name',)
admin.site.register(course_requirements,display_req)

class display_acccess_to(admin.ModelAdmin):
    list_display = ('course_name','access_to')
    list_per_page = 10
    search_fields = ('course_name','access_to')
    list_filter = ('course_name',)
admin.site.register(access_to,display_acccess_to)

admin.site.register(course_comment)
admin.site.register(course_comment_reply)

class display_enroll(admin.ModelAdmin):
    list_display = ('user_email','course_id','status')
    list_per_page = 10
    search_fields = ('user_email',)
admin.site.register(enrollment,display_enroll)


import sqlite3
import sys
import os
import csv
import math
import datetime
import debugMessage as dm
import numpy as np
import matplotlib.pyplot as plt
from os import walk
from utilities import *
from averageCSVFiles import *
#---------------------------------------------------------------------#
# This python file takes in multiple runs and calculate how Average   #
# Marginal Regret  changes as the app user percentage changes and the #
# simulation time elapses. The main execution is at the bottom of this#
# file.                                                               #
# NOTE: All "ND" refer to "Nash Distance," which is the old name of   #
#       average marginal regret. This will be fixed in later changes  #
#---------------------------------------------------------------------#

def generatePlot(percentage, timeStep, absNash, outputDir, plotEvent, eventStart, eventEnd, NDType = 'abs'):
  #-------------------------------------------------------------------#
  # percentage: Sorted list containing all app user percentages       #
  # timeStep: Time step for grouping vehicles                         #
  # absNash: Absolute Nash Distance                                   #
  # plotEven: Set true to plot two dash lines that mark the start and #
  #           end of events, specified by eventStart and eventEnd     #
  #-------------------------------------------------------------------#
  font = {'weight':'bold',
          'size':24}
  plt.rc('font', **font)
  fig, ax = plt.subplots(figsize = (24, 14), dpi = 100)

  #-------------------------------------------------------------------#
  # Change timestep into minutes from seconds                         #
  #-------------------------------------------------------------------#
  for i in range(len(timeStep)):
    timeStep[i] /= 60

  for i in range(len(percentage)):
    if NDType == 'abs':
      ax.plot(timeStep, absNash[i], color = (1 - float(percentage[i])/100, 0, float(percentage[i])/100), \
            label = '{}% App Users'.format(percentage[i]), \
            linewidth = 4.0)
    else:
      ax.plot(timeStep, absNash[i], color = (0, 1 - float(percentage[i])/100, float(percentage[i])/100), \
            label = '{}% App Users'.format(percentage[i]), \
            linewidth = 4.0)
    hand, lab = ax.get_legend_handles_labels()
    plt.legend(handles = hand, loc = 2)
  ax.set_xlabel('Elapsed Time (min)')
  if NDType == 'abs':
    ax.set_ylabel('Absolute Average Marginal Regret \n(sec)')
    ax.set_title('Absolute Average Marginal Regret for \nDifferent Percentages of App Users')
  else:
    ax.set_ylabel('Relative Average Marginal Regret \n(%)')
    ax.set_title('Relative Average Marginal Regret for \nDifferent Percentages of App Users')
  #-------------------------------------------------------------------#
  # Uncomment next line if you have to fix the upper/lower bounds of  #
  # the y-axis                                                        #
  #-------------------------------------------------------------------#
  #ax.set_xlim([0, 400])

  if plotEvent:
    ymin, ymax = plt.ylim()
    ax.plot([eventStart, eventStart], [ymin, ymax], color = (0, 0, 0), dashes = [2, 2], linewidth = 4.0)
    ax.plot([eventEnd, eventEnd], [ymin, ymax], color = (0, 0, 0), dashes = [2, 2], linewidth = 4.0)
  if NDType == 'abs':
    plt.savefig(outputDir + 'absND.png', dpi = 'figure')
  else:
    plt.savefig(outputDir + 'rltND.png', dpi = 'figure')

def extractSingleDB(fileName, thisPercentage, oriId, desId, maxT, minT, stepSize):
  #-------------------------------------------------------------------#
  # Extraction of a single SQLite database and calculate how absolute #
  # Nash Distance changes over time                                   #
  #                                                                   #
  # maxTime/minTime: Max/Min allowed time of entrance for vehicles    #
  # odiId/desId: id of desired origin/destination                     #
  #-------------------------------------------------------------------#
  print(('extracting database ' + fileName))

  con = sqlite3.connect(fileName)
  cur = con.cursor()
  cur.execute('SELECT oid, entranceTime, (exitTime - entranceTime) FROM MIVEHTRAJECTORY WHERE exitTime != -1 AND entranceTime > {} AND entranceTime < {} AND origin = {} AND destination = {}'.format(minT, maxT, oriId, desId))
  vehTTime = cur.fetchall()
  con.close()

  tStepAccumTime = {}
  tStepAccumCount= {}
  tStepMinTime   = {}
  for i in range(len(vehTTime)):
    thisTStep = math.floor((vehTTime[i][1] - minT)/stepSize)
    if thisTStep in tStepAccumTime:
      tStepAccumTime[thisTStep] += vehTTime[i][2]
      tStepAccumCount[thisTStep]+= 1
      if vehTTime[i][2] < tStepMinTime[thisTStep]:
        tStepMinTime[thisTStep] = vehTTime[i][2]
    else:
      tStepAccumTime[thisTStep]  = vehTTime[i][2]
      tStepAccumCount[thisTStep] = 1
      tStepMinTime[thisTStep]    = vehTTime[i][2]
  tStep   = []
  absNash = []
  rltNash = []
  for i in range(int(math.ceil((maxT - minT) / stepSize))):
    tStep.append(i * stepSize)
    if i in tStepMinTime:
      absNash.append(tStepAccumTime[i]/tStepAccumCount[i] - tStepMinTime[i])
      rltNash.append(100 * absNash[-1] * tStepAccumCount[i] / tStepAccumTime[i])
    else:
      absNash.append(0)
      rltNash.append(0)
  return tStep, absNash, rltNash

def sortBasedOnPercentage(percentage, absNash, rltNash):
  #-------------------------------------------------------------------#
  # Helper function to sort all data based on the percentage of app   #
  # users                                                             #
  #-------------------------------------------------------------------#
  zipped = list(zip(percentage, absNash, rltNash))
  zipped.sort()
  unzip  = list(zip(*zipped))
  sPercentage, sAbsNash, sRltNash = list(map(list, unzip))
  return sPercentage, sAbsNash, sRltNash

def traverseMultiDB(fileList, maxTime, minTime, stepSize):
  #-------------------------------------------------------------------#
  # Traverse multiple SQLite DB and return how Absolute Nash Distance #
  # changes with time                                                 #
  #-------------------------------------------------------------------#
  percentage = []
  absNash    = []
  rltNash    = []
  tStep      = []
  oriId, desId  = getODPair(fileList[0])
  for filename in fileList:
    thisPercentage  = getPercentage(filename)
    thisTStep, thisAbsNash, thisRltNash = extractSingleDB(filename, thisPercentage, oriId, desId, maxTime, minTime, stepSize)
    percentage.append(thisPercentage)
    absNash.append(thisAbsNash)
    rltNash.append(thisRltNash)
    tStep = thisTStep
  percentage, absNash, rltNash = sortBasedOnPercentage(percentage, absNash, rltNash)
  return percentage, tStep, absNash, rltNash

def printUsage():
  print('usage: \n\t python std4_timeND.py directoryName maxEntranceTime minEntranceTime stepSize')
  print('\t directoryName: the directory in which the sqlite databases are stored')
  print('\t stepSize: the time step size for grouping vehicles, in seconds')
  print('try \n\t python std4_timeND.py ../sqliteDatabases/DemoBenchmarkAccident/ 14400 7200 600')
  sys.exit()

def store2CSV(percentage, absND, runIdx, outputDir):
  outname = outputDir + 'run' + str(runIdx) + '_timeND.csv'
  with open(outname, 'w') as fout:
    for i in range(len(percentage)):
      fout.write(str(percentage[i]))
      for j in range(len(absND[i])):
        fout.write(',')
        fout.write(str(absND[i][j]))
      if i != (len(percentage) - 1):
        fout.write('\n')

# Main code starts here
if __name__ == '__main__':
  if len(sys.argv) != 5:
    printUsage()
  #-------------------------------------------------------------------#
  # Parsing input arguments and building list of all directories in   #
  # target (mother) directory                                         #
  #-------------------------------------------------------------------#
  dirList, outputDir, maxTime, minTime, stepSize = parseArgvTime(sys.argv, "timeAbsND")
  _, outputDirRlt, _, _, _ = parseArgvTime(sys.argv, "timeRltND")

  #-------------------------------------------------------------------#
  # Traverse, extract, and write data                                 #
  #-------------------------------------------------------------------#
  for dirName in dirList:
    runIdx   = getRunIdx(dirName)
    fileList = getAllFilenames(dirName)
    percentage, timeStep, absNash, rltNash = traverseMultiDB(fileList, maxTime, minTime, stepSize)
    store2CSV(percentage, absNash, runIdx, outputDir)
    store2CSV(percentage, rltNash, runIdx, outputDirRlt)

  #-------------------------------------------------------------------#
  # Read data from csv file and plot                                  #
  #-------------------------------------------------------------------#
  percentage, timeStep, absNash = averageTimeND(outputDir, stepSize)
  generatePlot(percentage, timeStep, absNash, '../outputFigures/', True, 2*60, 4*60)
  percentage, timeStep, rltNash = averageTimeND(outputDirRlt, stepSize)
  generatePlot(percentage, timeStep, rltNash, '../outputFigures/', True, 2*60, 4*60, 'rlt')


import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import torch
from monai.networks.blocks.activation import MemoryEfficientSwish, Swish

device = "cuda" if torch.cuda.is_available() else "cpu"

shapes_to_test = [128 * (x + 1) for x in range(8)]

def SwishTest(num_runs=1, mem_eff_swish=False):
    shapes_tested = []
    memory_used_in_mb = []

    if mem_eff_swish:
        layer = MemoryEfficientSwish()
    else:
        layer = Swish()

    for image_size in shapes_to_test:
        mem = 0.0
        for r in range(num_runs):
            input_image = torch.rand((1, 3, image_size, image_size), requires_grad=True).to(device)
            pred = layer(input_image)

            mem += torch.cuda.memory_allocated() / 1024 ** 2

            pred.sum().backward()

        memory_used_in_mb.append(mem / float(num_runs))
        shapes_tested.append(image_size)

    return shapes_tested, memory_used_in_mb

if __name__ == "__main__":
    sizes_tested_eff, memory_used_eff = SwishTest(mem_eff_swish=True)
    sizes_tested_neff, memory_used_neff = SwishTest(mem_eff_swish=False)

    plt.plot(sizes_tested_eff, memory_used_eff, "g-o", label="Memory Efficient Swish")
    plt.plot(sizes_tested_neff, memory_used_neff, "r-o", label="Swish")
    plt.xticks(ticks=sizes_tested_eff, rotation=45)
    plt.xlabel("Layer size")
    plt.ylabel("Allocated memory (MB)")
    plt.legend()
    plt.tight_layout()
    plt.grid()
    plt.savefig("figures/figure_singlelayer.png")


import socket
from tkinter import *
root = Tk()
sock = socket.socket()
sock.bind( ('', 9090) )
sock.listen(1)
print('waiting for connection...')
conn, addr = sock.accept()
print('connected: ', addr)
data = conn.recv(1024)
uData = data.decode("utf-8")
print("Client.py > " + uData)
conn.send(uData.encode("utf-8"))
conn.close()


from unittest import mock

from oslo_i18n import fixture as oslo_i18n_fixture
from oslotest import base as test_base

from oslo_utils import encodeutils

class EncodeUtilsTest(test_base.BaseTestCase):

    def test_safe_decode(self):
        safe_decode = encodeutils.safe_decode
        self.assertRaises(TypeError, safe_decode, True)
        self.assertEqual('ni\xf1o',
                         safe_decode("ni\xc3\xb1o".encode("latin-1"),
                                     incoming="utf-8"))

        self.assertEqual("strange",
                         safe_decode('\x80strange'.encode("latin-1"),
                                     errors='ignore'))

        self.assertEqual('\xc0', safe_decode('\xc0'.encode("latin-1"),
                         incoming='iso-8859-1'))

        # Forcing incoming to ascii so it falls back to utf-8
        self.assertEqual('ni\xf1o',
                         safe_decode('ni\xc3\xb1o'.encode("latin-1"),
                                     incoming='ascii'))

        self.assertEqual('foo', safe_decode(b'foo'))

    def test_safe_encode_none_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, None)

    def test_safe_encode_bool_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, True)

    def test_safe_encode_int_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, 1)

    def test_safe_encode_list_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, [])

    def test_safe_encode_dict_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, {})

    def test_safe_encode_tuple_instead_of_text(self):
        self.assertRaises(TypeError, encodeutils.safe_encode, ('foo', 'bar', ))

    def test_safe_encode_force_incoming_utf8_to_ascii(self):
        # Forcing incoming to ascii so it falls back to utf-8
        self.assertEqual(
            'ni\xc3\xb1o'.encode("latin-1"),
            encodeutils.safe_encode('ni\xc3\xb1o'.encode("latin-1"),
                                    incoming='ascii'),
        )

    def test_safe_encode_same_encoding_different_cases(self):
        with mock.patch.object(encodeutils, 'safe_decode', mock.Mock()):
            utf8 = encodeutils.safe_encode(
                'foo\xf1bar', encoding='utf-8')
            self.assertEqual(
                encodeutils.safe_encode(utf8, 'UTF-8', 'utf-8'),
                encodeutils.safe_encode(utf8, 'utf-8', 'UTF-8'),
            )
            self.assertEqual(
                encodeutils.safe_encode(utf8, 'UTF-8', 'utf-8'),
                encodeutils.safe_encode(utf8, 'utf-8', 'utf-8'),
            )
            encodeutils.safe_decode.assert_has_calls([])

    def test_safe_encode_different_encodings(self):
        text = 'foo\xc3\xb1bar'
        result = encodeutils.safe_encode(
            text=text, incoming='utf-8', encoding='iso-8859-1')
        self.assertNotEqual(text, result)
        self.assertNotEqual("foo\xf1bar".encode("latin-1"), result)

    def test_to_utf8(self):
        self.assertEqual(encodeutils.to_utf8(b'a\xe9\xff'),        # bytes
                         b'a\xe9\xff')
        self.assertEqual(encodeutils.to_utf8(u'a\xe9\xff\u20ac'),  # Unicode
                         b'a\xc3\xa9\xc3\xbf\xe2\x82\xac')
        self.assertRaises(TypeError, encodeutils.to_utf8, 123)     # invalid

        # oslo.i18n Message objects should also be accepted for convenience.
        # It works because Message is a subclass of str. Use the
        # lazy translation to get a Message instance of oslo_i18n.
        msg = oslo_i18n_fixture.Translation().lazy("test")
        self.assertEqual(encodeutils.to_utf8(msg),
                         b'test')

class ExceptionToUnicodeTest(test_base.BaseTestCase):

    def test_str_exception(self):
        # The regular Exception class cannot be used directly:
        # Exception(u'\xe9').__str__() raises an UnicodeEncodeError
        # on Python 2
        class StrException(Exception):
            def __init__(self, value):
                Exception.__init__(self)
                self.value = value

            def __str__(self):
                return self.value

        # On Python 3, an exception which returns bytes with is __str__()
        # method (like StrException(bytes)) is probably a bug, but it was not
        # harder to support this silly case in exception_to_unicode().

        # Decode from ASCII
        exc = StrException(b'bytes ascii')
        self.assertEqual(encodeutils.exception_to_unicode(exc),
                         u'bytes ascii')

        # Decode from UTF-8
        exc = StrException(b'utf-8 \xc3\xa9\xe2\x82\xac')
        self.assertEqual(encodeutils.exception_to_unicode(exc),
                         u'utf-8 \xe9\u20ac')

        # Force the locale encoding to ASCII to test the fallback
        with mock.patch.object(encodeutils, '_getfilesystemencoding',
                               return_value='ascii'):
            # Fallback: decode from ISO-8859-1
            exc = StrException(b'rawbytes \x80\xff')
            self.assertEqual(encodeutils.exception_to_unicode(exc),
                             u'rawbytes \x80\xff')

        # No conversion needed
        exc = StrException(u'unicode ascii')
        self.assertEqual(encodeutils.exception_to_unicode(exc),
                         u'unicode ascii')

        # No conversion needed
        exc = StrException(u'unicode \xe9\u20ac')
        self.assertEqual(encodeutils.exception_to_unicode(exc),
                         u'unicode \xe9\u20ac')

        # Test the locale encoding
        with mock.patch.object(encodeutils, '_getfilesystemencoding',
                               return_value='koi8_r'):
            exc = StrException(b'\xf2\xd5\xd3\xd3\xcb\xc9\xca')
            # Decode from the locale encoding
            # (the message cannot be decoded from ASCII nor UTF-8)
            self.assertEqual(encodeutils.exception_to_unicode(exc),
                             u'\u0420\u0443\u0441\u0441\u043a\u0438\u0439')

    def test_oslo_i18n_message(self):
        # use the lazy translation to get a Message instance of oslo_i18n
        exc = oslo_i18n_fixture.Translation().lazy("test")
        self.assertEqual(encodeutils.exception_to_unicode(exc),
                         u"test")


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

from collections import defaultdict
from collections import OrderedDict
import logging
import os

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import json_tricks as json
import numpy as np

from dataset.JointsDataset import JointsDataset
from nms.nms import oks_nms
from nms.nms import soft_oks_nms

import time

logger = logging.getLogger(__name__)

class DeepFashion2Dataset(JointsDataset):
    '''
    "gt_class_keypoints_dict":{
        1: (0, 25),
        2: (25, 58),
        3: (58, 89),
        4: (89, 128),
        5: (128, 143),
        6: (143, 158),
        7: (158, 168),
        8: (168, 182),
        9: (182, 190),
        10: (190, 219),
        11: (219, 256),
        12: (256, 275),
        13: (275, 294)
    },
    '''
    def __init__(self, cfg, root, image_set, is_train, transform=None):
        super().__init__(cfg, root, image_set, is_train, transform)
        self.nms_thre = cfg.TEST.NMS_THRE
        self.image_thre = cfg.TEST.IMAGE_THRE
        self.soft_nms = cfg.TEST.SOFT_NMS
        self.oks_thre = cfg.TEST.OKS_THRE
        self.in_vis_thre = cfg.TEST.IN_VIS_THRE
        self.bbox_file = cfg.TEST.DEEPFASHION2_BBOX_FILE
        # self.bbox_file = cfg.TEST.COCO_BBOX_FILE
        self.use_gt_bbox = cfg.TEST.USE_GT_BBOX
        self.image_width = cfg.MODEL.IMAGE_SIZE[0]
        self.image_height = cfg.MODEL.IMAGE_SIZE[1]
        self.mini_dataset = cfg.DATASET.MINI_DATASET
        self.select_cat = cfg.DATASET.SELECT_CAT
        self.aspect_ratio = self.image_width * 1.0 / self.image_height
        self.pixel_std = 200

        self.coco = COCO(self._get_ann_file_keypoint())

        # deal with class names
        cats = [cat['name']
                for cat in self.coco.loadCats(self.coco.getCatIds())]
        self.classes = ['__background__'] + cats
        logger.info('=> classes: {}'.format(self.classes))
        self.num_classes = len(self.classes)
        self._class_to_ind = dict(zip(self.classes, range(self.num_classes)))
        self._class_to_coco_ind = dict(zip(cats, self.coco.getCatIds()))
        self._coco_ind_to_class_ind = dict(
            [
                (self._class_to_coco_ind[cls], self._class_to_ind[cls])
                for cls in self.classes[1:]
            ]
        )

        # load image file names
        self.image_set_index = self._load_image_set_index()
        self.num_images = len(self.image_set_index)
        logger.info('=> num_images: {}'.format(self.num_images))

        self.num_joints = 294
        self.gt_class_keypoints_dict = {1: (0, 25), 2: (25, 58), 3: (58, 89),
                4: (89, 128), 5: (128, 143), 6: (143, 158), 7: (158, 168),
                8: (168, 182), 9: (182, 190), 10: (190, 219),
                11: (219, 256), 12: (256, 275), 13: (275, 294)}

        _flip_pairs = [
            [[2,6],[3,5],[7,25],[8,24],[9,23],[10,22],[11,21],[12,20],[13,19],[14,18],[15,17]],
            [[2,6],[3,5],[7,33],[8,32],[9,31],[10,30],[11,29],[12,28],[13,27],[14,26],[15,25],[16,24],[17,23],[18,22],[19,21]],
            [[2,26],[3,5],[4,6],[7,25],[8,24],[9,23],[10,22],[11,21],[12,20],[13,19],[14,18],[15,17],[16,29],[27,30],[28,31]],
            [[2,6],[3,5],[4,34],[7,33],[8,32],[9,31],[10,30],[11,29],[12,28],[13,27],[14,26],[15,25],[16,24],[17,23],[18,22],[19,21],[20,37],[35,38],[36,39]],
            [[2,6],[3,5],[7,15],[8,14],[9,13],[10,12]],
            [[2,6],[3,5],[7,15],[8,14],[9,13],[10,12]],
            [[1,3],[4,10],[5,9],[6,8]],
            [[1,3],[4,14],[5,13],[6,12],[7,11],[8,10]],
            [[1,3],[4,8],[5,7]],
            [[2,6],[3,5],[7,29],[8,28],[9,27],[10,26],[11,25],[12,24],[13,23],[14,22],[15,21],[16,20],[17,19]],
            [[2,6],[3,5],[7,37],[8,36],[9,35],[10,34],[11,33],[12,32],[13,31],[14,30],[15,29],[16,28],[17,27],[18,26],[19,25],[20,24],[21,23]],
            [[2,6],[3,5],[7,19],[8,18],[9,17],[10,16],[11,15],[12,14]],
            [[2,6],[3,5],[7,19],[8,18],[9,17],[10,16],[11,15],[12,14]]
        ]
        self.flip_pairs = []
        for idx, _cat_pairs in enumerate(_flip_pairs):
            start_idx = self.gt_class_keypoints_dict[idx+1][0]
            cat_pairs = []
            for pair in _cat_pairs:
                x0 = pair[0] + start_idx - 1
                x1 = pair[1] + start_idx - 1
                cat_pairs.append([x0,x1])
            self.flip_pairs.append(cat_pairs)

        self.parent_ids = None

        self.joints_weight = np.ones((self.num_joints, 1), dtype=np.float32
                                    ).reshape((self.num_joints, 1))

        print('Generating samples...')
        tic = time.time()
        self.cls_stat = np.array([0 for i in range(self.num_classes)])
        self.sample_list_of_cls = []
        self.db = self._get_db()
        print('Done (t={:0.2f}s)'.format(time.time()- tic))

        if is_train and cfg.DATASET.SELECT_DATA:
            self.db = self.select_data(self.db)

        logger.info('=> load {} samples'.format(len(self.db)))

    def _get_ann_file_keypoint(self):
        if 'train'  in self.image_set:
            directory = self.image_set
            if self.mini_dataset:
                return os.path.join(self.root, directory, 'train-coco_style-32.json')
            else:
                return os.path.join(self.root, directory, 'train-coco_style.json')
        elif 'validation'  in self.image_set:
            directory = self.image_set
            if self.mini_dataset:
                return os.path.join(self.root, directory, 'val-coco_style-64.json')
            else:
                return os.path.join(self.root, directory, 'val-coco_style.json')
        elif 'test'  in self.image_set:
            directory = 'json_for_test'
            return os.path.join(self.root, directory, 'keypoints_test_information.json')
        else:
            raise NotImplementedError

    def _load_image_set_index(self):
        """ image id: int """
        image_ids = self.coco.getImgIds()
        return image_ids

    def _get_db(self):
        if self.is_train or self.use_gt_bbox:
            # use ground truth bbox
            gt_db = self._load_coco_keypoint_annotations()
        else:
            # use bbox from detection
            # gt_db = self._load_coco_person_detection_results()
            gt_db = self._load_deepfashion2_detection_results()
        return gt_db

    def _load_coco_keypoint_annotations(self):
        """ ground truth bbox and keypoints """
        gt_db = []
        for index in self.image_set_index:
            gt_db.extend(self._load_coco_keypoint_annotation_kernal(index))
        return gt_db

    def _load_coco_keypoint_annotation_kernal(self, index):
        """
        coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
        iscrowd:
            crowd instances are handled by marking their overlaps with all categories to -1
            and later excluded in training
        bbox:
            [x1, y1, w, h]
        :param index: coco image id
        :return: db entry
        """
        im_ann = self.coco.loadImgs(index)[0]
        width = im_ann['width']
        height = im_ann['height']

        annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
        objs = self.coco.loadAnns(annIds)

        # sanitize bboxes
        valid_objs = []
        for obj in objs:
            x, y, w, h = obj['bbox']
            x1 = np.max((0, x))
            y1 = np.max((0, y))
            x2 = np.min((width - 1, x1 + np.max((0, w - 1))))
            y2 = np.min((height - 1, y1 + np.max((0, h - 1))))
            if obj['area'] > 0 and x2 >= x1 and y2 >= y1:
                obj['clean_bbox'] = [x1, y1, x2-x1, y2-y1]
                valid_objs.append(obj)
        objs = valid_objs

        rec = []
        for obj in objs:
            cls = self._coco_ind_to_class_ind[obj['category_id']]

            # ignore objs of not chosen class, or without keypoints annotation
            if cls not in self.select_cat or max(obj['keypoints']) == 0:
                continue
            self.cls_stat[cls] += 1
            self.sample_list_of_cls.append(cls)

            joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
            joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float)
            for ipt in range(self.num_joints):
                joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0]
                joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1]
                joints_3d[ipt, 2] = 0
                t_vis = obj['keypoints'][ipt * 3 + 2]
                if t_vis > 1:
                    t_vis = 1
                joints_3d_vis[ipt, 0] = t_vis
                joints_3d_vis[ipt, 1] = t_vis
                joints_3d_vis[ipt, 2] = 0

            center, scale = self._box2cs(obj['clean_bbox'][:4])
            rec.append({
                'image': self.image_path_from_index(index),
                'center': center,
                'scale': scale,
                'joints_3d': joints_3d,
                'joints_3d_vis': joints_3d_vis,
                'filename': '',
                'imgnum': 0,
                'category_id':obj['category_id'],
                'area': obj['area']
            })
        return rec

    def _box2cs(self, box):
        x, y, w, h = box[:4]
        return self._xywh2cs(x, y, w, h)

    def _xywh2cs(self, x, y, w, h):
        center = np.zeros((2), dtype=np.float32)
        center[0] = x + w * 0.5
        center[1] = y + h * 0.5

        if w > self.aspect_ratio * h:
            h = w * 1.0 / self.aspect_ratio
        elif w < self.aspect_ratio * h:
            w = h * self.aspect_ratio
        scale = np.array(
            [w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std],
            dtype=np.float32)
        if center[0] != -1:
            scale = scale * 1.25

        return center, scale

    def image_path_from_index(self, index):
        file_name = '%06d.jpg' % index
        file_name = os.path.join('image', file_name)

        image_path = os.path.join(self.root, self.image_set, file_name)

        return image_path

    def _load_deepfashion2_detection_results(self):
        all_boxes = None
        bbox_file_type = self.bbox_file.split('.')[-1]

        if bbox_file_type == 'json':
            with open(self.bbox_file, 'r') as f:
                all_boxes = json.load(f)

            if not all_boxes:
                logger.error('=> Load %s fail!' % self.bbox_file)
                return None

            logger.info('=> Total boxes: {}'.format(len(all_boxes)))

            kpt_db = []
            num_boxes = 0
            for n_img in range(0, len(all_boxes)):
                det_res = all_boxes[n_img]
                # if det_res['category_id'] != 1:
                #     continue
                img_name = self.image_path_from_index(det_res['image_id'])
                box = det_res['bbox']
                score = det_res['score']

                if score < self.image_thre:
                    continue

                num_boxes = num_boxes + 1

                center, scale = self._box2cs(box)
                joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
                joints_3d_vis = np.ones(
                    (self.num_joints, 3), dtype=np.float)
                kpt_db.append({
                    'image': img_name,
                    'center': center,
                    'scale': scale,
                    'score': score,
                    'joints_3d': joints_3d,
                    'joints_3d_vis': joints_3d_vis,
                })
        elif bbox_file_type == 'pkl':
            logger.info("Loading detection results from %s ..." % self.bbox_file)
            import pickle
            with open(self.bbox_file, 'rb') as f:
                raw_data = pickle.load(f)
                # all_boxes = self._process_pickle(raw_data)
                all_boxes = self._process_pickle(raw_data)
            if not all_boxes:
                logger.error('=> Load %s fail!' % self.bbox_file)
                return None

            logger.info('=> Total boxes: {}'.format(len(all_boxes)))

            kpt_db = []
            num_boxes = 0
            for n_img in range(0, len(all_boxes)):
                det_res = all_boxes[n_img]
                # if det_res['category_id'] != 1:
                #     continue
                img_name = self.image_path_from_index(det_res['image_id'])
                box = det_res['bbox']
                score = det_res['score']
                category_id = det_res['category_id']

                if score < self.image_thre:
                    continue

                num_boxes = num_boxes + 1

                center, scale = self._box2cs(box)
                joints_3d = np.zeros((self.num_joints, 3), dtype=np.float)
                joints_3d_vis = np.ones(
                    (self.num_joints, 3), dtype=np.float)
                kpt_db.append({
                    'image': img_name,
                    'center': center,
                    'scale': scale,
                    'score': score,
                    'joints_3d': joints_3d,
                    'joints_3d_vis': joints_3d_vis,
                    'category_id': det_res['category_id']
                })

        logger.info('=> Total boxes after fliter low score@{}: {}'.format(
            self.image_thre, num_boxes))
        return kpt_db

    def _process_pickle(self, data):
        all_boxes = []
        for n_img in range(len(data)):
            for i in range(13):
                for entry in data[n_img][0][i]:
                    # print('img:%07d cat:%02d' % (n_img, i+1))
                    # input()
                    if entry is not None:
                        x1, y1, x2, y2 = entry[0:4]
                        w = x2 - x1
                        h = y2 - y1
                        box = dict()
                        box['image_id'] = n_img + 1
                        box['bbox'] = [x1, y1, w, h]
                        box['score'] = entry[4]
                        box['category_id'] = i + 1
                    all_boxes.append(box)
        return all_boxes

    def evaluate(self, cfg, preds, output_dir, all_boxes, img_path,
                 *args, **kwargs):
        rank = cfg.RANK

        res_folder = os.path.join(output_dir, 'results')
        if not os.path.exists(res_folder):
            try:
                os.makedirs(res_folder)
            except Exception:
                logger.error('Fail to make {}'.format(res_folder))

        res_file = os.path.join(
            res_folder, 'keypoints_{}_results_{}.json'.format(
                self.image_set, rank)
        )

        # person x (keypoints)
        _kpts = []
        for idx, kpt in enumerate(preds):
            _kpts.append({
                'keypoints': kpt,
                'center': all_boxes[idx][0:2],
                'scale': all_boxes[idx][2:4],
                'area': all_boxes[idx][4],
                'score': all_boxes[idx][5],
                'category_id': all_boxes[idx][6],
                'image': int(img_path[idx][-10:-4])
            })
        # image x person x (keypoints)
        kpts = defaultdict(list)
        for kpt in _kpts:
            kpts[kpt['image']].append(kpt)

        # rescoring and oks nms
        num_joints = self.num_joints
        in_vis_thre = self.in_vis_thre
        oks_thre = self.oks_thre
        oks_nmsed_kpts = []
        for img in kpts.keys():
            img_kpts = kpts[img]
            for n_p in img_kpts:
                box_score = n_p['score']
                kpt_score = 0
                valid_num = 0
                for n_jt in range(0, num_joints):
                    t_s = n_p['keypoints'][n_jt][2]
                    if t_s > in_vis_thre:
                        kpt_score = kpt_score + t_s
                        valid_num = valid_num + 1
                if valid_num != 0:
                    kpt_score = kpt_score / valid_num
                # rescoring
                n_p['score'] = kpt_score * box_score
                # n_p['score'] = box_score

            if self.soft_nms:
                keep = soft_oks_nms(
                    [img_kpts[i] for i in range(len(img_kpts))],
                    oks_thre
                )
            else:
                keep = oks_nms(
                    [img_kpts[i] for i in range(len(img_kpts))],
                    oks_thre
                )

            if len(keep) == 0:
                oks_nmsed_kpts.append(img_kpts)
            else:
                oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep])

        # self._write_coco_keypoint_results(oks_nmsed_kpts, res_file)
        self._write_coco_keypoint_results_DeepFashion2(oks_nmsed_kpts, res_file)
        if 'test' not in self.image_set:
            info_str = self._do_python_keypoint_eval(
                res_file, res_folder)
            name_value = OrderedDict(info_str)
            return name_value, name_value['AP']
        else:
            return {'Null': 0}, 0

    def _write_coco_keypoint_results_DeepFashion2(self, keypoints, res_file):
        results = self._coco_keypoint_results_all_category_kernel(keypoints)
        logger.info('=> writing results json to %s' % res_file)
        with open(res_file, 'w') as f:
            json.dump(results, f, sort_keys=True, indent=4)
        try:
            json.load(open(res_file))
        except Exception:
            content = []
            with open(res_file, 'r') as f:
                for line in f:
                    content.append(line)
            content[-1] = ']'
            with open(res_file, 'w') as f:
                for c in content:
                    f.write(c)

    def _coco_keypoint_results_all_category_kernel(self, keypoints):
        # cat_id = data_pack['cat_id']

        keypoints = keypoints
        cat_results = []

        for img_kpts in keypoints:
            if len(img_kpts) == 0:
                continue

            for k in range(len(img_kpts)):
                image_id = img_kpts[k]['image']
                cat_id = img_kpts[k]['category_id']
                score = img_kpts[k]['score']

                _key_points = np.array([img_kpts[k]['keypoints']
                                        for k in range(len(img_kpts))])
                key_points = np.zeros(
                    (_key_points.shape[0], self.num_joints * 3), dtype=np.float
                )

                for ipt in range(self.num_joints):
                    key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0]
                    key_points[:, ipt * 3 + 1] = _key_points[:, ipt, 1]
                    key_points[:, ipt * 3 + 2] = _key_points[:, ipt, 2]  # keypoints score.

                result = [
                    {
                        'image_id': image_id,
                        'category_id': cat_id,
                        'keypoints': list(key_points[k]),
                        'score': score,
                        # 'score': img_kpts[k]['score'],
                        # 'center': list(img_kpts[k]['center']),
                        # 'scale': list(img_kpts[k]['scale'])
                    }
                ]
                cat_results.extend(result)

        return cat_results

    def _do_python_keypoint_eval(self, res_file, res_folder):
        coco_dt = self.coco.loadRes(res_file)
        coco_eval = COCOeval(self.coco, coco_dt, 'keypoints')
        coco_eval.params.useSegm = None
        coco_eval.evaluate()
        coco_eval.accumulate()
        coco_eval.summarize()

        stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)']

        info_str = []
        for ind, name in enumerate(stats_names):
            info_str.append((name, coco_eval.stats[ind]))

        return info_str

    def get_channel_index(self, class_id):
        gt_class_keypoints_dict =  {
            1: (0, 25), 2: (25, 58), 3: (58, 89), 4: (89, 128), 5: (128, 143),
            6: (143, 158), 7: (158, 168), 8: (168, 182), 9: (182, 190),
            10: (190, 219), 11: (219, 256), 12: (256, 275), 13: (275, 294)}

        if isinstance(class_id, int):
            return list(range(gt_class_keypoints_dict[class_id]))
        elif isinstance(class_id, list):
            return [list(range(gt_class_keypoints_dict[i])) for i in class_id]

    def generate_target(self, joints, joints_vis):
        '''
        :param joints:  [num_joints, 3]
        :param joints_vis: [num_joints, 3]
        :return: target, target_weight(1: visible, 0: invisible)
        '''
        # target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
        target_weight = np.zeros((self.num_joints, 1), dtype=np.float32)
        target_weight[:, 0] = joints_vis[:, 0]
        feat_stride = self.image_size / self.heatmap_size

        if self.target_type == 'gaussian':
            target = np.zeros((self.num_joints,
                               self.heatmap_size[1],
                               self.heatmap_size[0]),
                              dtype=np.float32)

            tmp_size = self.sigma * 3

            for joint_id in range(self.num_joints):
                mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
                mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
                # Check that any part of the gaussian is in-bounds
                ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
                br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
                if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
                        or br[0] < 0 or br[1] < 0:
                    # If not, just return the image as is
                    target_weight[joint_id] = 0
                    continue

                # # Generate gaussian
                size = 2 * tmp_size + 1
                x = np.arange(0, size, 1, np.float32)
                y = x[:, np.newaxis]
                x0 = y0 = size // 2
                # The gaussian is not normalized, we want the center value to equal 1
                g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))

                # Usable gaussian range
                g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
                g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
                # Image range
                img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
                img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])

                v = target_weight[joint_id]
                if v > 0.5:
                    target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
                        g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
        elif self.target_type == 'coordinate':
            target = joints[:, 0:2]
            target /= feat_stride
        else:
            raise NotImplementedError('Only support gaussian map and coordinate now!')

        if self.use_different_joints_weight:
            target_weight = np.multiply(target_weight, self.joints_weight)

        return target, target_weight


import numpy as np
import time
import cv2

# using weights and configs from
# https://www.pyimagesearch.com/2018/11/12/yolo-object-detection-with-opencv/
class YoloV3Detector:
    def __init__(self):
        self.min_confidence = 0.5
        self.threshold = 0.3
        self.scale_factor = 1/255.0
        self.image_size = (416,416)
        self.detection_color = (0,255,0)
        self.tracking_color = (255,0,0)
        self.label_path = 'yolo/coco.names'
        self.weights_path = 'yolo/yolov3.weights'
        self.config_path = 'yolo/yolov3.cfg'
        self.labels = open(self.label_path).read().strip().split("\n")
        self.yolo_detector = cv2.dnn.readNetFromDarknet(self.config_path,
                self.weights_path)
        self.layer_names = self.yolo_detector.getLayerNames()
        # yolo3 has three different output layers - we can use all three
        self.output_layers = self.yolo_detector.getUnconnectedOutLayers()
        # these are all the yolo_layers (-1 bc i don't think the layers are returned in
        # 0 index fashion)
        self.output_layer_names = [self.layer_names[i[0]-1] for i in
                self.output_layers]
        self.image_shape = None
        self.outputs = None
        self.boxes = []
        self.confidences = []
        self.classIDs = []
        self.COLORS = None
        self.image = None

    def createBlobAndForwardPass(self, image):
        # create blob from image and perform a forward pass
        # blob performs mean subtraction & scaling for us see here
        # https://www.pyimagesearch.com/2017/11/06/deep-learning-opencvs-blobfromimage-works/
        self.image = image
        self.image_shape = image.shape[:2]
        blob = cv2.dnn.blobFromImage(image, self.scale_factor, self.image_size, swapRB=True,
                crop=False)
        self.yolo_detector.setInput(blob)
        self.outputs = self.yolo_detector.forward(self.output_layer_names)

    def createBoundingBoxes(self):
        # get image dimensions
        (h, w) = self.image_shape

        # reset bins
        self.boxes = []
        self.confidences = []
        self.classIDs = []

        for output in self.outputs:
            for detection in output:
                # each row is a candidate detection, the 1st 4 numbers are
                # [center_x, center_y, width, height], followed by (N-4) class
                # probabilities/confidence
                scores = detection[5:]
                classID = np.argmax(scores)
                confidence = scores[classID]

                # filter out weak predictions
                if confidence < self.min_confidence:
                    continue

                # for this assignment, only care about sports ball label
                if self.labels[classID] != 'sports ball':
                    continue

                # scale bounding box to size of image
                (center_x, center_y, width, height) = (detection[0:4] * np.array([w, h,
                    w, h])).astype('int')

                # grab the corners of the box
                left = int(center_x - width/2)
                top = int(center_y - height/2)

                # add new values to their arrays
                self.boxes.append([left, top, int(width), int(height)])
                self.confidences.append(float(confidence))
                self.classIDs.append(classID)

    def applyNonMaxSuppression(self):
        # apply non-maxima suppression to overlapping bounding boxes
        self.final_detections = cv2.dnn.NMSBoxes(self.boxes, self.confidences,
                self.min_confidence, self.threshold)

    def drawBoundingBoxesWithLabels(self):
        if len(self.final_detections) > 0:
            boxes = self.boxes
            for d in self.final_detections.flatten():
                # draw bounding box
                color = [int(c) for c in self.COLORS[self.classIDs[d]]]
                self.drawBoundingBox(self.image, boxes[d], color)
                text = "{}: {:.4f}".format(self.labels[self.classIDs[d]], self.confidences[d])
                cv2.putText(self.image, text, (x,y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color,
                    1)

    def drawBoundingBox(self, image, bbox, color=None):
        color = color if color != None else self.detection_color
        print(bbox)
        (x, y) = (int(bbox[0]), int(bbox[1]))
        (w, h) = (int(bbox[2]), int(bbox[3]))
        # draw bounding box
        cv2.rectangle(image, (x,y), (x+w, y+h), color, 2)

    def returnSingleBoundingBox(self):
         if len(self.final_detections) > 0:
            boxes = self.boxes
            # this is based on the assumption that there is only a single sports
            # ball in the image and we care about nothing else
            singleBox = self.final_detections.flatten()[0]
            print('singleBox', boxes[singleBox])
            return tuple(boxes[singleBox])
         else:
            # could not find the ball
            return None

    def showImage(self):
        cv2.imshow("image", self.image)
        # cv2.waitKey(0)
        cv2.waitKey(25)

    def showImageWithAnnotations(self, image, bbox, color, annotation1,
            annotation2):
        if bbox != None:
            self.drawBoundingBox(image, bbox, color)
            # Display tracker type on frame
            cv2.putText(image, annotation1, (20,20),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2);

            # Display FPS on frame
            cv2.putText(image, annotation2, (20,50),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.75, color, 2);
        else:
            # Tracking failure
            cv2.putText(image, "Tracking failure detected", (20,80),
                      cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)

        cv2.imshow("video", image)

    def createColors(self):
        np.random.seed(42) # keeps colors the same everytime
        self.COLORS = np.random.randint(0, 255, size=(len(self.labels), 3), dtype='uint8')

    def detectAndShow(self, image):
       self.createColors()
       self.createBlobAndForwardPass(image)
       self.createBoundingBoxes()
       self.applyNonMaxSuppression()
       self.drawBoundingBoxesWithLabels()
       self.showImage()

    def detectAndReturnBoundingBox(self, image):
       self.createColors()
       self.createBlobAndForwardPass(image)
       self.createBoundingBoxes()
       self.applyNonMaxSuppression()
       return self.returnSingleBoundingBox()

# load image for testing
image_path = './images/soccer.jpg'
image = cv2.imread(image_path)

# load video
videoPath = './videos/soccer_game.mp4'
video = cv2.VideoCapture(videoPath)

# Exit if video not opened.
if not video.isOpened():
    print("Could not open video")

# Read first frame.
ok, frame = video.read()
if not ok:
    print('Cannot read video file')

# initalize yolov3 detector
yolov3 = YoloV3Detector()
bbox = yolov3.detectAndReturnBoundingBox(frame)

count = 0
ok = True
fps = 0

while True:
    okay, frame = video.read()
    if not okay:
        break

    # Start timer
    timer = cv2.getTickCount()

    if count % 20 == 0 or not ok:
        method = 'detection'
        color = yolov3.detection_color
        #detection takes some time, so im showing that we're detecting here too
        yolov3.showImageWithAnnotations(frame, bbox, color, method, "FPS : " +
                    str(int(fps)))
        # detect sports ball
        bbox = yolov3.detectAndReturnBoundingBox(frame)
        # create KCF tracker
        tracker = cv2.TrackerKCF_create()
        tracker.init(frame, bbox)
        ok = bbox != None
    else:
        # Update tracker
        ok, bbox = tracker.update(frame)
        method = 'tracking'
        color = yolov3.tracking_color

    # Calculate processing time and display results.
    # Calculate Frames per second (FPS)
    fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer);
    yolov3.showImageWithAnnotations(frame, bbox, color, method, "FPS : " +
            str(int(fps)))

    cv2.waitKey(1)
    count += 1


from datetime import datetime

def allowed(options):
    value = options['value'];
    allowed = options['parameter']

    output = {}
    if value in allowed:
        output['success'] = True;
    else:
        output['success'] = False;
        output['message'] = value + ' not in (' + '|'.join(allowed) + ')'

    return output

def dateformat(options):
    formats = options['parameter']
    date = options['value']

    output = {}
    for format in formats:
        try:
            output['success'] = True;
        except ValueError:
            output['success'] = False;
            output['message'] = 'date did not match format (' + '|'.join(formats) + ')'

    return output;


# python
from decimal import Decimal
from model_bakery import baker

# django
from django.contrib.auth import get_user_model
from django.forms.models import model_to_dict
from django.urls import reverse

# third party
from rest_framework.test import APITestCase
from rest_framework import status
from rest_framework_simplejwt.tokens import RefreshToken

# project
from core.constants import DATETIME_FORMAT
from core.serializers import UserSerializer

# local
from loans.models import Loan
from loans.models import Payment
from loans.constants import PRICE_SYSTEM
from loans.constants import SAC_SYSTEM
from .utils import utc_to_local

User = get_user_model()

class BaseAPITestCase(APITestCase):

    def setUp(self):
        # admin
        self.admin = baker.make(User, is_staff=True, is_superuser=True)
        self.admin_token = RefreshToken.for_user(self.admin).access_token
        self.admin_headers = {'REMOTE_ADDR': '192.168.0.10'}

        # client
        self.user = baker.make(User)
        self.headers = {'REMOTE_ADDR': '192.168.0.20'}
        self.token = RefreshToken.for_user(self.user).access_token

        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.admin_token}')

class BaseLoanAPITestCase(BaseAPITestCase):

    def setUp(self):
        super().setUp()

        data_price = {'client': self.user,
                      'bank': 'testbank',
                      'value': 20000.00,
                      'interest_rate': 0.04,
                      'period': 8,
                      'financing': PRICE_SYSTEM}
        self.loan_price = Loan.objects.create(**data_price)

        user = baker.make(User)
        data_sac = {'client': user,
                    'bank': 'testbank',
                    'value': 120000.00,
                    'interest_rate': 0.05,
                    'period': 10,
                    'financing': SAC_SYSTEM}
        self.loan_sac = Loan.objects.create(**data_sac)

class TestLoanCreateAPIView(BaseAPITestCase):

    def get_url(self):
        return reverse('loans:create')

    def test_create_loan_without_authentication(self):
        self.client.logout()

        data = {'client': self.user.id,
                'bank': 'testbank',
                'value': 20000.00,
                'interest_rate': 0.04,
                'period': 8,
                'financing': PRICE_SYSTEM}
        response = self.client.post(self.get_url(), data, **self.admin_headers)

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'As credenciais de autenticação não foram fornecidas.')
        self.assertEqual(Loan.objects.count(), 0)
        self.assertEqual(Payment.objects.count(), 0)

    def test_create_loan_by_client(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')

        data = {'client': self.user.id,
                'bank': 'testbank',
                'value': 20000.00,
                'interest_rate': 0.04,
                'period': 8,
                'financing': PRICE_SYSTEM}
        response = self.client.post(self.get_url(), data, **self.headers)

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'Você não tem permissão para executar essa ação.')
        self.assertEqual(Loan.objects.count(), 0)
        self.assertEqual(Payment.objects.count(), 0)

    def test_create_loan_by_admin(self):
        data = {'client': self.user.id,
                'bank': 'testbank',
                'value': 20000.00,
                'interest_rate': 0.04,
                'period': 8,
                'financing': PRICE_SYSTEM}
        response = self.client.post(self.get_url(), data, **self.admin_headers)

        self.assertEqual(response.status_code, status.HTTP_201_CREATED)
        self.assertEqual(Loan.objects.count(), 1)
        self.assertEqual(Payment.objects.count(), 8)

        loan = Loan.objects.get()
        self.assertEqual(loan.ip_address, '192.168.0.10')
        self.assertEqual(loan.amount_due, Decimal('23764.48'))

    def test_create_loan_with_empty_data(self):
        reponse_error = {
            'client': ['Este campo é obrigatório.'],
            'bank': ['Este campo é obrigatório.'],
            'value': ['Este campo é obrigatório.'],
            'interest_rate': ['Este campo é obrigatório.'],
            'period': ['Este campo é obrigatório.']}

        data = {}
        response = self.client.post(self.get_url(), data, **self.admin_headers)

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertDictEqual(response.json(), reponse_error)
        self.assertEqual(Loan.objects.count(), 0)
        self.assertEqual(Payment.objects.count(), 0)

    def test_create_loan_with_wrong_data(self):
        reponse_error = {
            'client': ['Pk inválido "10" - objeto não existe.'],
            'value': ['Certifque-se de que este valor seja maior ou igual a 0.01.'],
            'interest_rate': ['Certifque-se de que este valor seja maior ou igual a 0.01.'],
            'period': ['Um número inteiro válido é exigido.'],
            'financing': ['"0" não é um escolha válido.']}

        data = {'client': 10,
                'bank': 'testbank',
                'value': 0.00,
                'interest_rate': -0.01,
                'period': 'a',
                'financing': 0}
        response = self.client.post(self.get_url(), data, **self.admin_headers)

        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertDictEqual(response.json(), reponse_error)
        self.assertEqual(Loan.objects.count(), 0)
        self.assertEqual(Payment.objects.count(), 0)

class TestLoanListAPIView(BaseLoanAPITestCase):

    def get_url(self):
        return reverse('loans:list')

    def test_list_loan_without_authentication(self):
        self.client.logout()

        response = self.client.get(self.get_url())

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'As credenciais de autenticação não foram fornecidas.')

    def test_list_loan_by_client(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')

        response = self.client.get(self.get_url())

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.json().get('results')), 1)

    def test_list_loan_by_admin(self):
        response = self.client.get(self.get_url())

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.json().get('results')), 2)

class TestLoanRetrieveAPIView(BaseLoanAPITestCase):

    maxDiff = None

    def get_url(self, loan):
        return reverse('loans:retrieve', args=[loan.id])

    def serializer_format(self, loan):
        return {
            **model_to_dict(loan),
            'id': str(loan.id),
            'created': utc_to_local(loan.created).strftime(DATETIME_FORMAT),
            'modified': utc_to_local(loan.modified).strftime(DATETIME_FORMAT),
            'client': UserSerializer(loan.client).data,
            'value': f'R$ {loan.value:.2f}',
            'amount_due': f'R$ {loan.amount_due:.2f}',
            'interest_rate': f'{100.0 * float(loan.interest_rate):.2f}%',
            'balance_due': f'R$ {loan.make_balance_due():.2f}'}

    def test_retrieve_loan_price_without_authentication(self):
        self.client.logout()

        response = self.client.get(self.get_url(self.loan_price))

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'As credenciais de autenticação não foram fornecidas.')

    def test_retrieve_loan_price_by_client(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')
        response_json = self.serializer_format(self.loan_price)

        response = self.client.get(self.get_url(self.loan_price))

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertDictEqual(response.json(), response_json)

    def test_retrieve_loan_sac_by_client_without_permission(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')

        response = self.client.get(self.get_url(self.loan_sac))

        self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
        self.assertEqual(response.json().get('detail'), 'Não encontrado.')

    def test_retrieve_loan_price_by_admin(self):
        response = self.client.get(self.get_url(self.loan_price))
        response_json = self.serializer_format(self.loan_price)

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertDictEqual(response.json(), response_json)

    def test_retrieve_loan_sac_by_admin(self):
        response = self.client.get(self.get_url(self.loan_sac))
        response_json = self.serializer_format(self.loan_sac)

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertDictEqual(response.json(), response_json)

class TestPaymentListAPIView(BaseLoanAPITestCase):

    def get_url(self, loan):
        return reverse('loans:payments-list', args=[loan.id])

    def test_list_loan_price_payments_without_authentication(self):
        self.client.logout()

        response = self.client.get(self.get_url(self.loan_price))

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'As credenciais de autenticação não foram fornecidas.')

    def test_list_loan_price_payment_by_client(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')

        response = self.client.get(self.get_url(self.loan_price))

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.json().get('results')), 8)

    def test_list_loan_sac_payment_by_client_without_permission(self):
        self.client.credentials(HTTP_AUTHORIZATION=f'Bearer {self.token}')

        response = self.client.get(self.get_url(self.loan_sac))

        self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
        self.assertEqual(response.json().get('detail'), 'Você não tem permissão para executar essa ação.')

    def test_list_loan_price_payment_by_admin(self):
        response = self.client.get(self.get_url(self.loan_price))

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.json().get('results')), 8)

    def test_list_loan_sac_payment_by_admin(self):
        response = self.client.get(self.get_url(self.loan_sac))

        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(len(response.json().get('results')), 10)


import logging
from typing import Dict, List, Optional, Tuple

import reagent.types as rlt
import torch
from reagent.models.base import ModelBase
from reagent.models.seq2slate import Seq2SlateMode, Seq2SlateTransformerNet
from reagent.models.seq2slate_reward import Seq2SlateRewardNetBase
from reagent.preprocessing.postprocessor import Postprocessor
from reagent.preprocessing.preprocessor import Preprocessor
from reagent.preprocessing.sparse_preprocessor import (
    SparsePreprocessor,
    make_sparse_preprocessor,
)
from reagent.torch_utils import gather
from reagent.training.utils import gen_permutations
from torch import nn

logger = logging.getLogger(__name__)

def serving_to_feature_data(
    serving: rlt.ServingFeatureData,
    dense_preprocessor: Preprocessor,
    sparse_preprocessor: SparsePreprocessor,
) -> rlt.FeatureData:
    float_features_with_presence, id_list_features, id_score_list_features = serving
    return rlt.FeatureData(
        float_features=dense_preprocessor(*float_features_with_presence),
        id_list_features=sparse_preprocessor.preprocess_id_list(id_list_features),
        id_score_list_features=sparse_preprocessor.preprocess_id_score_list(
            id_score_list_features
        ),
    )

def sparse_input_prototype(
    model: ModelBase,
    state_preprocessor: Preprocessor,
    state_feature_config: rlt.ModelFeatureConfig,
):
    name2id = state_feature_config.name2id
    model_prototype = model.input_prototype()
    # Terrible hack to make JIT tracing works. Python dict doesn't have type
    # so we need to insert something so JIT tracer can infer the type.
    state_id_list_features = {
        42: (torch.zeros(1, dtype=torch.long), torch.tensor([], dtype=torch.long))
    }
    state_id_score_list_features = {
        42: (
            torch.zeros(1, dtype=torch.long),
            torch.tensor([], dtype=torch.long),
            torch.tensor([], dtype=torch.float),
        )
    }
    if isinstance(model_prototype, rlt.FeatureData):
        if model_prototype.id_list_features:
            state_id_list_features = {
                name2id[k]: v for k, v in model_prototype.id_list_features.items()
            }
        if model_prototype.id_score_list_features:
            state_id_score_list_features = {
                name2id[k]: v for k, v in model_prototype.id_score_list_features.items()
            }

    input = rlt.ServingFeatureData(
        float_features_with_presence=state_preprocessor.input_prototype(),
        id_list_features=state_id_list_features,
        id_score_list_features=state_id_score_list_features,
    )
    return (input,)

class DiscreteDqnWithPreprocessor(ModelBase):
    """
    This is separated from DiscreteDqnPredictorWrapper so that we can pass typed inputs
    into the model. This is possible because JIT only traces tensor operation.
    In contrast, JIT scripting needs to compile the code, therefore, it won't recognize
    any custom Python type.
    """

    def __init__(
        self,
        model: ModelBase,
        state_preprocessor: Preprocessor,
        state_feature_config: rlt.ModelFeatureConfig,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.state_feature_config = state_feature_config
        self.sparse_preprocessor = make_sparse_preprocessor(
            self.state_feature_config, device=torch.device("cpu")
        )

    def forward(self, state: rlt.ServingFeatureData):
        state_feature_data = serving_to_feature_data(
            state, self.state_preprocessor, self.sparse_preprocessor
        )
        q_values = self.model(state_feature_data)
        return q_values

    def input_prototype(self):
        return sparse_input_prototype(
            model=self.model,
            state_preprocessor=self.state_preprocessor,
            state_feature_config=self.state_feature_config,
        )

class DiscreteDqnPredictorWrapper(torch.jit.ScriptModule):
    def __init__(
        self,
        dqn_with_preprocessor: DiscreteDqnWithPreprocessor,
        action_names: List[str],
        # here to keep interface consistent with FB internal
        state_feature_config: rlt.ModelFeatureConfig,
    ) -> None:
        super().__init__()
        self.dqn_with_preprocessor = torch.jit.trace(
            dqn_with_preprocessor, dqn_with_preprocessor.input_prototype()
        )
        self.action_names = torch.jit.Attribute(action_names, List[str])

    # pyre-fixme[56]: Decorator `torch.jit.script_method` could not be resolved in a
    #  global scope.
    @torch.jit.script_method
    def forward(self, state: rlt.ServingFeatureData) -> Tuple[List[str], torch.Tensor]:
        q_values = self.dqn_with_preprocessor(state)
        return (self.action_names, q_values)

class OSSSparsePredictorUnwrapper(nn.Module):
    # Wrap input in serving feature data
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        state_id_list_features: Dict[int, Tuple[torch.Tensor, torch.Tensor]],
        state_id_score_list_features: Dict[
            int, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        ],
    ) -> Tuple[List[str], torch.Tensor]:
        return self.model(
            rlt.ServingFeatureData(
                float_features_with_presence=state_with_presence,
                id_list_features=state_id_list_features,
                id_score_list_features=state_id_score_list_features,
            )
        )

# Pass through serving module's output
class OSSPredictorUnwrapper(nn.Module):
    def __init__(self, model: nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(self, *args, **kwargs) -> Tuple[List[str], torch.Tensor]:
        return self.model(*args, **kwargs)

DiscreteDqnPredictorUnwrapper = OSSSparsePredictorUnwrapper
ActorPredictorUnwrapper = OSSPredictorUnwrapper
ParametricDqnPredictorUnwrapper = OSSPredictorUnwrapper

class ParametricDqnWithPreprocessor(ModelBase):
    def __init__(
        self,
        model: ModelBase,
        state_preprocessor: Preprocessor,
        action_preprocessor: Preprocessor,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.action_preprocessor = action_preprocessor

    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        action_with_presence: Tuple[torch.Tensor, torch.Tensor],
    ):
        preprocessed_state = self.state_preprocessor(
            state_with_presence[0], state_with_presence[1]
        )
        preprocessed_action = self.action_preprocessor(
            action_with_presence[0], action_with_presence[1]
        )
        state = rlt.FeatureData(preprocessed_state)
        action = rlt.FeatureData(preprocessed_action)
        q_value = self.model(state, action)
        return q_value

    def input_prototype(self):
        return (
            self.state_preprocessor.input_prototype(),
            self.action_preprocessor.input_prototype(),
        )

class ParametricDqnPredictorWrapper(torch.jit.ScriptModule):
    def __init__(self, dqn_with_preprocessor: ParametricDqnWithPreprocessor) -> None:
        super().__init__()

        self.dqn_with_preprocessor = torch.jit.trace(
            dqn_with_preprocessor, dqn_with_preprocessor.input_prototype()
        )

    # pyre-fixme[56]: Decorator `torch.jit.script_method` could not be resolved in a
    #  global scope.
    @torch.jit.script_method
    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        action_with_presence: Tuple[torch.Tensor, torch.Tensor],
    ) -> Tuple[List[str], torch.Tensor]:
        value = self.dqn_with_preprocessor(state_with_presence, action_with_presence)
        return (["Q"], value)

class ActorWithPreprocessor(ModelBase):
    """
    This is separate from ActorPredictorWrapper so that we can pass typed inputs
    into the model. This is possible because JIT only traces tensor operation.
    In contrast, JIT scripting needs to compile the code, therefore, it won't recognize
    any custom Python type.
    """

    def __init__(
        self,
        model: ModelBase,
        state_preprocessor: Preprocessor,
        action_postprocessor: Optional[Postprocessor] = None,
        serve_mean_policy: bool = False,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.action_postprocessor = action_postprocessor
        self.serve_mean_policy = serve_mean_policy

    def forward(self, state_with_presence: Tuple[torch.Tensor, torch.Tensor]):
        preprocessed_state = self.state_preprocessor(
            state_with_presence[0], state_with_presence[1]
        )
        state_feature_vector = rlt.FeatureData(preprocessed_state)
        # TODO: include log_prob in the output
        model_output = self.model(state_feature_vector)
        if self.serve_mean_policy:
            assert (
                model_output.squashed_mean is not None
            ), "action mean is None and serve_mean_policy=True"
            action = model_output.squashed_mean
        else:
            action = model_output.action

        if self.action_postprocessor:
            # pyre-fixme[29]: `Optional[Postprocessor]` is not a function.
            action = self.action_postprocessor(action)
        return action

    def input_prototype(self):
        return (self.state_preprocessor.input_prototype(),)

_DEFAULT_FEATURE_IDS = []

class ActorPredictorWrapper(torch.jit.ScriptModule):
    def __init__(
        self,
        actor_with_preprocessor: ActorWithPreprocessor,
        action_feature_ids: List[int] = _DEFAULT_FEATURE_IDS,
    ) -> None:
        """
        action_feature_ids is here to make the interface consistent with FB internal
        version
        """
        super().__init__()

        self.actor_with_preprocessor = torch.jit.trace(
            actor_with_preprocessor, actor_with_preprocessor.input_prototype()
        )

    # pyre-fixme[56]: Decorator `torch.jit.script_method` could not be resolved in a
    #  global scope.
    @torch.jit.script_method
    def forward(
        self, state_with_presence: Tuple[torch.Tensor, torch.Tensor]
    ) -> torch.Tensor:
        action = self.actor_with_preprocessor(state_with_presence)
        return action

class Seq2SlateWithPreprocessor(ModelBase):
    def __init__(
        self,
        model: Seq2SlateTransformerNet,
        state_preprocessor: Preprocessor,
        candidate_preprocessor: Preprocessor,
        greedy: bool,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.candidate_preprocessor = candidate_preprocessor
        self.greedy = greedy

    def input_prototype(self):
        candidate_input_prototype = self.candidate_preprocessor.input_prototype()
        return (
            self.state_preprocessor.input_prototype(),
            (
                candidate_input_prototype[0].repeat((1, self.model.max_src_seq_len, 1)),
                candidate_input_prototype[1].repeat((1, self.model.max_src_seq_len, 1)),
            ),
        )

    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        candidate_with_presence: Tuple[torch.Tensor, torch.Tensor],
    ):
        # state_value.shape == state_presence.shape == batch_size x state_feat_num
        # candidate_value.shape == candidate_presence.shape ==
        # batch_size x max_src_seq_len x candidate_feat_num
        batch_size = state_with_presence[0].shape[0]

        preprocessed_state = self.state_preprocessor(
            state_with_presence[0], state_with_presence[1]
        )
        preprocessed_candidates = self.candidate_preprocessor(
            candidate_with_presence[0].view(
                batch_size * self.model.max_src_seq_len,
                len(self.candidate_preprocessor.sorted_features),
            ),
            candidate_with_presence[1].view(
                batch_size * self.model.max_src_seq_len,
                len(self.candidate_preprocessor.sorted_features),
            ),
        ).view(batch_size, self.model.max_src_seq_len, -1)

        # TODO: consider different numbers of candidates in the same batch_
        src_src_mask = torch.ones(
            batch_size, self.model.max_src_seq_len, self.model.max_src_seq_len
        )
        ranking_input = rlt.PreprocessedRankingInput.from_tensors(
            state=preprocessed_state,
            src_seq=preprocessed_candidates,
            src_src_mask=src_src_mask,
        )
        ranking_output = self.model(
            ranking_input,
            mode=Seq2SlateMode.RANK_MODE,
            tgt_seq_len=self.model.max_tgt_seq_len,
            greedy=self.greedy,
        )
        return ranking_output.ranked_tgt_out_probs, ranking_output.ranked_tgt_out_idx

class Seq2SlatePredictorWrapper(torch.jit.ScriptModule):
    def __init__(self, seq2slate_with_preprocessor: Seq2SlateWithPreprocessor) -> None:
        super().__init__()
        self.seq2slate_with_preprocessor = torch.jit.trace(
            seq2slate_with_preprocessor, seq2slate_with_preprocessor.input_prototype()
        )

    # pyre-fixme[56]: Decorator `torch.jit.script_method` could not be resolved in a
    #  global scope.
    @torch.jit.script_method
    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        candidate_with_presence: Tuple[torch.Tensor, torch.Tensor],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # ranked_tgt_out_probs shape: batch_size, tgt_seq_len, candidate_size
        # ranked_tgt_out_idx shape: batch_size, tgt_seq_len
        ranked_tgt_out_probs, ranked_tgt_out_idx = self.seq2slate_with_preprocessor(
            state_with_presence, candidate_with_presence
        )
        # convert to slate-wise probabilities
        # ranked_tgt_out_probs shape: batch_size
        ranked_tgt_out_probs = torch.prod(
            torch.gather(
                ranked_tgt_out_probs, 2, ranked_tgt_out_idx.unsqueeze(-1)
            ).squeeze(),
            -1,
        )
        # -2 to offset padding symbol and decoder start symbol
        ranked_tgt_out_idx -= 2
        return ranked_tgt_out_probs, ranked_tgt_out_idx

class Seq2RewardWithPreprocessor(DiscreteDqnWithPreprocessor):
    def __init__(
        self,
        model: ModelBase,
        state_preprocessor: Preprocessor,
        seq_len: int,
        num_action: int,
    ):
        """
        Since TorchScript unable to trace control-flow, we
        have to generate the action enumerations as constants
        here so that trace can use them directly.
        """

        super().__init__(model, state_preprocessor, rlt.ModelFeatureConfig())
        self.seq_len = seq_len
        self.num_action = num_action
        self.all_permut = gen_permutations(seq_len, num_action)
        self.num_permut = self.all_permut.size(1)

    def forward(self, state: rlt.ServingFeatureData):
        """
        This serving module only takes in current state.
        We need to simulate all multi-step length action seq's
        then predict accumulated reward on all those seq's.
        After that, we categorize all action seq's by their
        first actions. Then take the maximum reward as the
        predicted categorical reward for that category.
        Return: categorical reward for the first action
        """
        state_with_presence, _, _ = state
        batch_size, state_dim = state_with_presence[0].size()

        # expand state tensor to match the enumerated action sequences:
        # the tensor manipulations here are tricky:
        # Suppose the input states are s1,s2, these manipulations
        # will generate a input batch s1,s1,...,s1,s2,s2,...,s2
        # where len(s1,s1,...,s1)=len(s2,s2,...,s2)=num_permut
        preprocessed_state = (
            self.state_preprocessor(state_with_presence[0], state_with_presence[1])
            .repeat(1, self.seq_len * self.num_permut)
            .reshape(batch_size * self.num_permut, self.seq_len, -1)
            .transpose(0, 1)
        )
        state_feature_vector = rlt.FeatureData(preprocessed_state)

        # expand action to match the expanded state sequence
        action = self.all_permut.repeat(1, batch_size, 1)
        reward = self.model(
            state_feature_vector, rlt.FeatureData(action)
        ).acc_reward.reshape(
            batch_size, self.num_action, self.num_permut // self.num_action
        )

        # The permuations are generated with lexical order
        # the output has shape [num_perm, num_action,1]
        # that means we can aggregate on the max reward
        # then reshape it to (BATCH_SIZE, ACT_DIM)
        max_reward = (
            # pyre-fixme[16]: `Tuple` has no attribute `values`.
            torch.max(reward, 2)
            .values.cpu()
            .detach()
            .reshape(batch_size, self.num_action)
        )

        return max_reward

class Seq2SlateRewardWithPreprocessor(ModelBase):
    def __init__(
        self,
        model: Seq2SlateRewardNetBase,
        state_preprocessor: Preprocessor,
        candidate_preprocessor: Preprocessor,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.candidate_preprocessor = candidate_preprocessor

    def input_prototype(self):
        candidate_input_prototype = self.candidate_preprocessor.input_prototype()

        return (
            self.state_preprocessor.input_prototype(),
            (
                candidate_input_prototype[0].repeat((1, self.model.max_src_seq_len, 1)),
                candidate_input_prototype[1].repeat((1, self.model.max_src_seq_len, 1)),
            ),
        )

    @property
    def state_sorted_features(self) -> List[int]:
        return self.state_preprocessor.sorted_features

    @property
    def candidate_sorted_features(self) -> List[int]:
        return self.candidate_preprocessor.sorted_features

    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        candidate_with_presence: Tuple[torch.Tensor, torch.Tensor],
    ):
        # state_value.shape == state_presence.shape == batch_size x state_feat_num
        # candidate_value.shape == candidate_presence.shape ==
        # batch_size x max_src_seq_len x candidate_feat_num
        batch_size = state_with_presence[0].shape[0]
        max_tgt_seq_len = self.model.max_tgt_seq_len
        max_src_seq_len = self.model.max_src_seq_len

        # we use a fake slate_idx_with_presence to retrive the first
        # max_tgt_seq_len candidates from
        # len(slate_idx_with presence) == batch_size
        # component: 1d tensor with length max_tgt_seq_len
        slate_idx_with_presence = [
            (torch.arange(max_tgt_seq_len), torch.ones(max_tgt_seq_len))
        ] * batch_size

        preprocessed_state = self.state_preprocessor(
            state_with_presence[0], state_with_presence[1]
        )

        preprocessed_candidates = self.candidate_preprocessor(
            candidate_with_presence[0].view(
                batch_size * max_src_seq_len, len(self.candidate_sorted_features)
            ),
            candidate_with_presence[1].view(
                batch_size * max_src_seq_len, len(self.candidate_sorted_features)
            ),
        ).view(batch_size, max_src_seq_len, -1)

        src_src_mask = torch.ones(batch_size, max_src_seq_len, max_src_seq_len)

        tgt_out_idx = torch.cat(
            [slate_idx[0] for slate_idx in slate_idx_with_presence]
        ).view(batch_size, max_tgt_seq_len)

        tgt_out_seq = gather(preprocessed_candidates, tgt_out_idx)

        ranking_input = rlt.PreprocessedRankingInput.from_tensors(
            state=preprocessed_state,
            src_seq=preprocessed_candidates,
            src_src_mask=src_src_mask,
            tgt_out_seq=tgt_out_seq,
            # +2 is needed to avoid two preserved symbols:
            # PADDING_SYMBOL = 0
            # DECODER_START_SYMBOL = 1
            tgt_out_idx=tgt_out_idx + 2,
        )

        output = self.model(ranking_input)
        return output.predicted_reward

class MDNRNNWithPreprocessor(ModelBase):
    def __init__(
        self,
        model: ModelBase,
        state_preprocessor: Preprocessor,
        seq_len: int,
        num_action: int,
        state_feature_config: Optional[rlt.ModelFeatureConfig] = None,
    ):
        super().__init__()
        self.model = model
        self.state_preprocessor = state_preprocessor
        self.state_feature_config = state_feature_config or rlt.ModelFeatureConfig()
        self.sparse_preprocessor = make_sparse_preprocessor(
            self.state_feature_config, device=torch.device("cpu")
        )
        self.seq_len = seq_len
        self.num_action = num_action

    def forward(
        self,
        state_with_presence: Tuple[torch.Tensor, torch.Tensor],
        action: torch.Tensor,
    ):

        batch_size, state_dim = state_with_presence[0].size()
        preprocessed_state = (
            self.state_preprocessor(state_with_presence[0], state_with_presence[1])
            .reshape(batch_size, self.seq_len, -1)
            .transpose(0, 1)
        )
        result = self.model(action, preprocessed_state)

        return result

    def input_prototype(self):
        return (
            self.state_preprocessor.input_prototype(),
            torch.randn(1, 1, self.num_action, device=self.state_preprocessor.device),
        )

class CompressModelWithPreprocessor(DiscreteDqnWithPreprocessor):
    def forward(self, state: rlt.ServingFeatureData):
        state_feature_data = serving_to_feature_data(
            state, self.state_preprocessor, self.sparse_preprocessor
        )
        # TODO: model is a fully connected network which only takes in Tensor now.
        q_values = self.model(state_feature_data.float_features)
        return q_values