import os
import six
import imageio
from PIL import Image
import numpy as np

from .images import ImagesDataset
from . import datasets_root
from . import util
from .utils import nputil
from .utils import imutil

try:
    # Python 2
    from itertools import izip
except ImportError:
    # Python 3
    izip = zip

from tensorflow.contrib.slim.nets import inception, vgg

class coco_animals(ImagesDataset):
    def __init__(self):
        ImagesDataset.__init__(self)
        self.dataset_name='coco-animals'
        self.source_url='http://cs231n.stanford.edu/'
        self.source_files=['coco-animals.zip']
        self.dataset_home=os.path.join(datasets_root,self.dataset_name)
        self.height=None
        self.width=None
        self.depth=None

        self.x_layout =  imutil.LAYOUT_NCHW
        self.x_layout_file = imutil.LAYOUT_NCHW

        #self.x_shape = 'NCHW' ## other alternates are NHW and NHCW
        self.n_classes = 8

    def load_data(self,force=False, shuffle=True, x_is_images=False):
        self.dataset_home=os.path.join(datasets_root,self.dataset_name)

        self.downloaded_files=util.download_dataset(source_url=self.source_url,
                                                    source_files=self.source_files,
                                                    dest_dir = self.dataset_home,
                                                    force=force,
                                                    extract=True)

        # the archive contains the name of the coco-animals folder, hence this temporary fix for now
        self.dataset_home=os.path.join(self.dataset_home,self.dataset_name)
#        print('Extracting ',self.downloaded_files[0])
        ilabels = os.listdir(os.path.join(self.dataset_home,'train'))

        label2id = dict(zip(ilabels, range(len(ilabels))))
        #print(label2id)

        self.label2id=label2id
        id2label = dict(zip(label2id.values(), label2id.keys()))
        self.id2label = id2label

        n_train = 800
        n_valid = 200

        x_train_files=[]
        y_train = np.zeros((n_train,), dtype=np.uint8)
        x_valid_files=[]
        y_valid=np.zeros((n_valid,), dtype=np.uint8)
        #print(ilabels)

        for i in range(self.n_classes):
            label = id2label[i]
            ifolder = os.path.join(self.dataset_home,'train',label)
            #print(os.listdir(ifolder))
            files = [name for name in os.listdir(ifolder) if name.endswith('.jpg')]
            for f in files:
                x_train_files.append(os.path.join(ifolder,f))
                #y_train.append(labels2id[label])
            y_train[i * (n_train // self.n_classes): (i+1) * (n_train // self.n_classes)] = i

            ifolder = os.path.join(self.dataset_home,'val',label)
            files = [name for name in os.listdir(ifolder) if name.endswith('.jpg')]
            for f in files:
                x_valid_files.append(os.path.join(ifolder,f))
                #y_val.append(labels2id[label])
            y_valid[i * (n_valid // self.n_classes): (i+1) * (n_valid // self.n_classes)] = i

        if shuffle:
            x_train_files, y_train = self.shuffle_xy(x_train_files,y_train)

        if x_is_images:
            x_train = self.load_images(x_train_files)
            x_valid = self.load_images(x_valid_files)
        else:
            x_train = x_train_files
            x_valid = x_train_files
        self.x_is_images=x_is_images

        self.part['X_train']=x_train
        self.part['Y_train']=y_train

        self.part['X_valid']=x_valid
        self.part['Y_valid']=y_valid

        return x_train, y_train, x_valid, y_valid

    def preprocess_for_vgg(self,incoming):
        if isinstance(incoming, six.string_types):
            img = self.load_image(incoming)
        elif isinstance(incoming, imageio.core.util.Image ):
            img = Image.fromarray(incoming)
        else: #
            img=incoming

        #print(type(img))

        img_size = vgg.vgg_16.default_image_size

        height = img_size
        width = img_size

        img = self.resize_image(img,height,width)
        img = self.pil_to_nparray(img)
        if len(img.shape)==2:   # greyscale or no channels then add three channels
            h=img.shape[0]
            w=img.shape[1]
            img = np.dstack([img]*3)

        means = np.array([[[123.68, 116.78, 103.94]]]) #shape=[1, 1, 3]
        try:
            img = img - means
        except Exception as ex:
            print('Error preprocessing ',incoming)
            print(ex)

        return img

    def preprocess_for_inception(self,incoming):

        img_size = inception.inception_v3.default_image_size

        height = img_size
        width = img_size

        if isinstance(incoming, six.string_types):
            img = self.load_image(incoming)
        elif isinstance(incoming, imageio.core.util.Image ):
            img = Image.fromarray(incoming)
        else: #
            img=incoming

        #print(type(img))

        img = self.resize_image(img,height,width)
        img = self.pil_to_nparray(img)
        if len(img.shape)==2:   # greyscale or no channels then add three channels
            h=img.shape[0]
            w=img.shape[1]
            img = np.dstack([img]*3)

        img = ((img/255.0) - 0.5) * 2.0

        return img


"""Serializers for ecommerce"""
from rest_framework import serializers

from applications.models import BootcampApplication
from ecommerce.models import Order, Line, Receipt
from klasses.serializers import BootcampRunSerializer, InstallmentSerializer

class PaymentSerializer(serializers.Serializer):
    """
    Serializer for payment API, used to do basic validation.
    """

    payment_amount = serializers.DecimalField(
        max_digits=20, decimal_places=2, min_value=0.01
    )
    application_id = serializers.IntegerField()

class OrderPartialSerializer(serializers.ModelSerializer):
    """
    Serializer for Order
    """

    class Meta:
        model = Order
        fields = ("id", "status", "created_on", "updated_on")

class ReceiptSerializer(serializers.ModelSerializer):
    """Serializer for receipts for a user"""

    class Meta:
        model = Receipt
        fields = ["id", "payment_method"]

class ApplicationOrderSerializer(serializers.ModelSerializer):
    """Serializer for orders that are part of a bootcamp application"""

    total_price_paid = serializers.DecimalField(decimal_places=2, max_digits=20)
    payment_method = serializers.SerializerMethodField()

    def get_payment_method(self, order):
        """Get the payment method used in the last receipt for the order"""
        if order.payment_type == Order.CYBERSOURCE_TYPE:
            # There should only be one receipt for an order most of the time, but it's possible
            # there is a duplicate or a Cybersource error in one of the receipts.
            receipt = order.receipt_set.order_by("id").last()
            return receipt.payment_method if receipt is not None else None
        elif order.payment_type == Order.WIRE_TRANSFER_TYPE:
            return "Wire Transfer"

    class Meta:
        model = Order
        fields = [
            "id",
            "status",
            "total_price_paid",
            "created_on",
            "updated_on",
            "payment_method",
        ]

class OrderSerializer(serializers.ModelSerializer):
    """Serializer for orders"""

    class Meta:
        model = Order
        fields = ["id", "status", "application_id", "created_on", "updated_on"]

class LineSerializer(serializers.ModelSerializer):
    """
    Serializer for Line
    """

    order = OrderPartialSerializer(read_only=True)
    price = serializers.DecimalField(decimal_places=2, max_digits=20)
    run_key = serializers.SerializerMethodField()

    def get_run_key(self, line):
        """get run_key from bootcamp_run"""
        return line.bootcamp_run.run_key

    class Meta:
        model = Line
        fields = ("order", "price", "description", "run_key")

class CheckoutDataSerializer(serializers.ModelSerializer):
    """Serializer for ecommerce information for a BootcampApplication"""

    bootcamp_run = BootcampRunSerializer()
    total_price = serializers.DecimalField(
        max_digits=20, decimal_places=2, read_only=True, source="price"
    )
    total_paid = serializers.DecimalField(
        max_digits=20, decimal_places=2, read_only=True
    )
    payments = serializers.SerializerMethodField()
    installments = serializers.SerializerMethodField()

    def get_payments(self, application):
        """Serialized payments made by the user"""
        return LineSerializer(
            (
                order.line_set.first()
                for order in application.orders.all()
                if order.status == Order.FULFILLED
            ),
            many=True,
        ).data

    def get_installments(self, application):
        """Installments with prices and due dates"""
        return InstallmentSerializer(
            application.bootcamp_run.installment_set.order_by("deadline"), many=True
        ).data

    class Meta:
        model = BootcampApplication
        fields = [
            "id",
            "bootcamp_run",
            "total_price",
            "total_paid",
            "payments",
            "installments",
        ]


import experiments

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import os

# lookups
default_split = {
    'cub200': 2,
    'cars196': 0,
    'sop': 0,
}

dataset_label_lookup = {
    'cub200': 'CUB200-2011',
    'cars196': 'CARS196',
    'sop': 'Online Products (SOP)'
}

# get wandb run data
name_list, group_list, summary_list, config_list, data_list = experiments.get_wandb_data()

project_group = 'ooDML_msloss'
labelspace = 'arch'
dataset = 'sop' # 'cub200'
filter_substrings = ['100epoch', 'vitS8Dino_bs30_c'] # 'vitS8Dino_bs30_c', 'gpu2_sop'

# filter runs data - project group
ids_valid = [i for i, group in enumerate(group_list) if group == project_group]
name_list = [name_list[i] for i in ids_valid]
data_list = [data_list[i] for i in ids_valid]

# filter runs data - dataset
ids_valid = [i for i, d in enumerate(data_list) if d['dataset'] == dataset]
name_list = [name_list[i] for i in ids_valid]
data_list = [data_list[i] for i in ids_valid]

# filter runs data - substrings in name
ids_valid = [i for i, n in enumerate(name_list) if not any(map(n.__contains__, filter_substrings))]
name_list = [name_list[i] for i in ids_valid]
data_list = [data_list[i] for i in ids_valid]

# get labels
labels_list = [d[labelspace] for d in data_list]
labels_unique = sorted(list(set(labels_list)))

# gather ooDML progressions
recall_seqs = {label: [] for label in labels_unique}
split_seqs = {label: [] for label in labels_unique}

for data in data_list:
    label_tmp = data[labelspace]
    recall_seqs[label_tmp].append(data['r@1_max'])
    split_seqs[label_tmp].append(data['ooDML_split_id'])

# sort progressions by split id
n_x_ticks = []
for label in labels_unique:
    recs = recall_seqs[label]
    split_ids = split_seqs[label]

    # sort recall values
    recs = [x for _, x in sorted(zip(split_ids, recs))]

    # reassign
    recall_seqs[label] = recs
    split_seqs[label] = sorted(split_ids)

    n_x_ticks.append(len(recs))
n_x_ticks = min(n_x_ticks)

# cut data to n_x_ticks - to still be able to visualize experiments in progress
for label in labels_unique:
    recall_seqs[label] = recall_seqs[label][:n_x_ticks]
    split_seqs[label] = split_seqs[label][:n_x_ticks]

# plotting
x_label = 'split_id'
x_ticks = list(range(1, n_x_ticks + 1)) # e.g. split ids
colors = ['green', 'red', 'blue', 'magenta', 'yellow', 'cyan', 'navy', 'orange', 'deeppink', 'springgreen', 'pink', 'darkkhaki']
alphas = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

# fig, ax = plt.subplots(figsize=(5, 4))
# for i, label in enumerate(labels_unique):
#     linestyle = 'dashed' if any(map(label.__contains__, ['resnet', 'bninception'])) else 'solid'
#     ax.plot(x_ticks, recall_seqs[label], label=f'{label}', marker='x', color=colors[i], linestyle=linestyle, alpha=alphas[i])
#
# leg = ax.legend(loc="best", shadow=True, fancybox=True, fontsize='xx-small')
# ax.set(xlabel=f'{x_label}', ylabel=f'{dataset} recall@1',
#        title=f'Eval ooDML [{labelspace}] [384]')
# ax.set_xticks(x_ticks)
# plt.axvline(x=x_ticks[default_split[dataset]]) # add default split indicator (vertical line)
# ax.yaxis.set_major_locator(ticker.MultipleLocator(0.02))
# ax.grid()

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
# absolute performance plot
for i, label in enumerate(labels_unique):
    linestyle = 'dashed' if any(map(label.__contains__, ['resnet', 'bninception'])) else 'solid'
    ax1.plot(x_ticks, recall_seqs[label], label=f'{label}', marker='x', color=colors[i], linestyle=linestyle, alpha=alphas[i])

leg = ax1.legend(loc="best", shadow=True, fancybox=True, fontsize='xx-small')
ax1.set(xlabel=f'{x_label}', ylabel=f'{dataset_label_lookup[dataset]} recall@1',
       title=f'Absolute recall performance [{labelspace}] [384]')
ax1.set_xticks(x_ticks)
plt.axvline(x=x_ticks[default_split[dataset]]) # add default split indicator (vertical line)
ax1.yaxis.set_major_locator(ticker.MultipleLocator(0.02))
ax1.grid()

# relative performance plot (reference value: performance for split 1)
for i, label in enumerate(labels_unique):
    linestyle = 'dashed' if any(map(label.__contains__, ['resnet', 'bninception'])) else 'solid'
    ax2.plot(x_ticks, [recall_seqs[label][0] - v for v in recall_seqs[label]], label=f'{label}', marker='x', color=colors[i], linestyle=linestyle, alpha=alphas[i])

leg = ax2.legend(loc="best", shadow=True, fancybox=True, fontsize='xx-small')
ax2.set(xlabel=f'{x_label}', ylabel=f'{dataset_label_lookup[dataset]} recall@1 difference',
       title=f'Relative recall performance w.r.t. split 1 [{labelspace}] [384]')
ax2.set_xticks(x_ticks)
plt.axvline(x=x_ticks[default_split[dataset]]) # add default split indicator (vertical line)
ax2.yaxis.set_major_locator(ticker.MultipleLocator(0.02))
ax2.grid()

# save plot
customstring = f'msloss_dino60_abs_rel'
path_save = f'/export/home/tmilbich/PycharmProjects/dml_pl/experiments/plots/ooDML/eval_{labelspace}_{dataset}'
path_save = "_".join([path_save, customstring])
fig.savefig(path_save + '.png', dpi=600) # .svg'

plt.show()


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import migrations, models

class Migration(migrations.Migration):

    dependencies = [
        ('fieldsight', '0002_auto_20161108_0034'),
        ('logger', '0005_remove_xform_site'),
    ]

    operations = [
        migrations.CreateModel(
            name='FieldSightXF',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('is_staged', models.BooleanField(default=False)),
                ('is_scheduled', models.BooleanField(default=False)),
                ('date_created', models.DateTimeField(auto_now=True)),
                ('date_modified', models.DateTimeField(auto_now=True)),
            ],
            options={
                'ordering': ('date_modified',),
                'db_table': 'fieldsight_forms_data',
                'verbose_name': 'XForm',
                'verbose_name_plural': 'XForms',
            },
        ),
        migrations.CreateModel(
            name='Schedule',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('name', models.CharField(max_length=256, verbose_name=b'Schedule Name')),
                ('date_range_start', models.DateField(auto_now=True)),
                ('date_range_end', models.DateField(auto_now=True)),
            ],
            options={
                'ordering': ('date_range_start',),
                'db_table': 'fieldsight_forms_schedule',
                'verbose_name': 'Form Schedule',
                'verbose_name_plural': 'Form Schedules',
            },
        ),
        migrations.CreateModel(
            name='Stage',
            fields=[
                ('id', models.AutoField(verbose_name='ID', serialize=False, auto_created=True, primary_key=True)),
                ('name', models.CharField(max_length=256)),
                ('order', models.IntegerField(default=0)),
                ('stage', models.ForeignKey(related_name='parent', blank=True, to='fsforms.Stage', null=True)),
            ],
            options={
                'ordering': ('order',),
                'db_table': 'fieldsight_forms_stage',
                'verbose_name': 'FieldSight Form Stage',
                'verbose_name_plural': 'FieldSight Form Stages',
            },
        ),
        migrations.AddField(
            model_name='fieldsightxf',
            name='schedule',
            field=models.ForeignKey(blank=True, to='fsforms.Schedule', null=True),
        ),
        migrations.AddField(
            model_name='fieldsightxf',
            name='site',
            field=models.ManyToManyField(related_name='site_forms', to='fieldsight.Site'),
        ),
        migrations.AddField(
            model_name='fieldsightxf',
            name='stage',
            field=models.ForeignKey(blank=True, to='fsforms.Stage', null=True),
        ),
        migrations.AddField(
            model_name='fieldsightxf',
            name='xf',
            field=models.ForeignKey(to='logger.XForm'),
        ),
    ]


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import re
import sys
import xml.etree.ElementTree as ElementTree
from pathlib import Path

FREE_MIND_NODE_TAG = 'node'
FREE_MIND_NODE_TEXT_ATTRIBUTE_TAG = 'TEXT'

time_re = re.compile(r'.+(\s+\(.*\))')
time_number_re = re.compile(r'.*\((\d+[.,]?\d?).*\)')

class NodeInfo:
    def __init__(self, text=""):
        self.text = text
        self.childs: [NodeInfo] = []
        self.time = None
        self.child_time = None

    def append_child(self, child):
        self.childs.append(child)

    def has_childs(self):
        return len(self.childs) > 0

    def update_time(self):

        if not self.has_childs():
            self.child_time = self.time
            return

        for child in self.childs:
            child.update_time()

            child_time = child.time if child.child_time is None else child.child_time

            if child_time is None:
                continue

            self.child_time = child_time if self.child_time is None else self.child_time + child_time

    def get_display_time(self):
        if self.time is None:
            return "" if self.child_time is None else str(self.child_time)

        if self.child_time is None:
            return "" if self.time is None else str(self.time)

        if self.time == self.child_time:
            return str(self.time)

        return f"{self.time} ?? {self.child_time}"

    def get_time(self):
        if self.time is None:
            return 0 if self.child_time is None else self.child_time

        if self.child_time is None:
            return 0 if self.time is None else self.time

        if self.time == self.child_time:
            return self.time

        return max(self.time, self.child_time)

def append_nodes(xml_parent_node, parent_node_info: NodeInfo):
    for child in xml_parent_node:
        if child.tag != FREE_MIND_NODE_TAG:
            continue

        if FREE_MIND_NODE_TEXT_ATTRIBUTE_TAG in child.attrib:
            node = NodeInfo(child.attrib[FREE_MIND_NODE_TEXT_ATTRIBUTE_TAG])

            time_match = time_re.match(node.text)

            if time_match is not None:
                time = time_match[1]
                node.text = node.text.replace(time, "")
                time_number_match = time_number_re.match(time)

                if time_number_match is not None:
                    node.time = float(time_number_match[1].replace(",", "."))

            parent_node_info.append_child(node)
            append_nodes(child, node)

def get_nodes_information(path: Path):
    tree = ElementTree.parse(path)
    root = tree.getroot()
    root_node_info = NodeInfo()

    for child in root:
        if child.tag != FREE_MIND_NODE_TAG:
            continue

        append_nodes(child, root_node_info)

    return root_node_info

def write_to_file_csv(outfile, node_info: NodeInfo, intend: str or None, number: str or None):
    if intend is None:
        node_intend = ""
        intend = ""
    else:
        node_intend = intend
        intend += "  "

    node_time = node_info.get_display_time().replace('.', ',')
    node_number = "" if not number else f"{number}."
    node_text = "Итого;" if not node_info.text else f"{node_info.text};"
    line = f"{node_intend}{node_number};{node_text}{node_time}"

    if line:
        outfile.write(line)
        outfile.write("\n")

    child_number = 1
    for child in node_info.childs:
        write_to_file_csv(outfile, child, intend, f"{node_number}{child_number}")
        child_number += 1

def write_to_file_skype_format(outfile, node_info: NodeInfo, intend: str or None, number: str or None, write_start: bool):
    if write_start:
        outfile.write("```\n")

    if intend is None:
        node_intend = ""
        intend = ""
    else:
        node_intend = intend
        intend += "  "

    node_time = node_info.get_display_time().replace(',', '.')
    node_time = node_time if not node_time else f" ({node_time} ч)"
    node_number = "" if not number else f"{number}."
    node_text = "Итого:" if not node_info.text else f" {node_info.text}"
    line = f"{node_intend}{node_number}{node_text}{node_time}"

    if line:
        outfile.write(line)
        outfile.write("\n")

    child_number = 1
    for child in node_info.childs:
        write_to_file_skype_format(outfile, child, intend, f"{node_number}{child_number}", False)
        child_number += 1

    if write_start:
        outfile.write("```\n")

def write_to_file_in_redmine_format_table(outfile, node_info: NodeInfo, number: str or None, write_start: bool):
    if write_start:
        outfile.write("|_. № п/п |_. Задача |_. Оценка, в часах) |\n")

    node_number = "" if not number else f"{number}."
    node_time = node_info.get_display_time().replace(',', '.')
    node_time = node_time if not node_time else f"{node_time}"

    if node_info.text:
        line = f"|{node_number}|{node_info.text}|{node_time}|"
    else:
        line = f"|*Итого:*| |_ *{node_time}*|"

    if line:
        outfile.write(line)
        outfile.write("\n")

    child_number = 1
    for child in node_info.childs:
        write_to_file_in_redmine_format_table(outfile, child, f"{node_number}{child_number}", False)
        child_number += 1

def write_to_file_in_redmine_format(outfile, node_info: NodeInfo, intend: str or None):
    if intend is None:
        node_intend = ""
        intend = "#"
    else:
        node_intend = intend
        intend += "#"

    node_time = node_info.get_display_time().replace(',', '.')
    node_time = node_time if not node_time else f"{node_time}"

    if node_info.text:
        line = f"{node_intend} {node_info.text}"

        if node_time:
            line += f" ({node_time})"
    else:
        line = f"*Итого: {node_time}*"

    if line:
        outfile.write(line)
        outfile.write("\n")

    child_number = 1
    for child in node_info.childs:
        write_to_file_in_redmine_format(outfile, child, intend)
        child_number += 1

def write_delimiter(outfile):
    outfile.write("\n-----\n")

def process_free_mind_document(path):
    in_file_path = Path(path)
    root_node_info = get_nodes_information(in_file_path)
    root_node_info.update_time()

    out_file_path = Path(in_file_path.parent).joinpath(in_file_path.name).with_suffix('.txt')

    if out_file_path.exists():
        out_file_path.unlink()

    with out_file_path.open('wt', encoding="utf-8") as outfile:
        outfile.write(f"{out_file_path.with_suffix('').name}\n")
        write_delimiter(outfile)
        write_to_file_csv(outfile, root_node_info, None, None)
        write_delimiter(outfile)
        write_to_file_in_redmine_format(outfile, root_node_info, None)
        write_delimiter(outfile)
        write_to_file_in_redmine_format_table(outfile, root_node_info, None, True)
        write_delimiter(outfile)
        write_to_file_skype_format(outfile, root_node_info, None, None, True)

if __name__ == "__main__":
    if len(sys.argv) < 2 or (sys.argv[1] is None):
        print('Usage: python3 convert-mm-to-excel.py path')
        sys.exit(-1)

    argv_iterator = iter(sys.argv)
    next(argv_iterator)

    for arg in argv_iterator:
        process_free_mind_document(arg)


import pika
import sys

connection = pika.BlockingConnection(
        pika.ConnectionParameters('localhost'))
channel = connection.channel()

# durable: 声明队列为持久化队列
channel.queue_declare(queue='task_queue', durable=True)

message = ' '.join(sys.argv[1:]) or "Hello World"
channel.basic_publish(
        exchange='',
        routing_key='task_queue',
        body=message,
        # delivery_mode=2: 将消息设置为持久化消息
        properties=pika.BasicProperties(
            delivery_mode=2,)  # make message persistent
        )
print("[x] Sent 'Hello World!'")
connection.close()


from urllib.request import urlopen
from bs4 import BeautifulSoup
import random

html = urlopen("https://dhlottery.co.kr/gameResult.do?method=statByNumber")
bsObject = BeautifulSoup(html, "html.parser")

theory_rate = 7/45 # 이론상 숫자가 나올 확률
current_round = int(bsObject.body.find("select",{"id":"edDrwNo"}).find("option",{"selected":""}).text) # 총 시행횟수 (최신 회차)
theory_pop_count = int(current_round * theory_rate) # 이론상 숫자가 나와야 하는 횟수

data = bsObject.body.find_all("tbody")[-1].text.split()
Pop_count = {} # 실제 숫자가 나온 횟수를 담을 딕셔너리
num = 1

for i in range(2,len(data),3):
    Pop_count[ num ] = int(data[i])
    num += 1

sorted_Pop = sorted(Pop_count.items(),key=(lambda x:x[1]), reverse=True)

most_high = sorted_Pop[:6]
most_low = sorted_Pop[-6:]
middle = sorted_Pop[ (int(len(sorted_Pop)/2)) - 3 : (int(len(sorted_Pop)/2)) + 3 ]

print("most high :", most_high)
print("most low :", most_low)

random_seq = most_high + most_low
random.shuffle(random_seq)

random_seq = random_seq[:6]
random_seq.sort()

print("shuffled :", random_seq)


from django.db.models import fields
from rest_framework import serializers
from .models import Order, OrderRecipe, Recipe
from recipe.serializers import RecipeSerializer

class OrderRecipeSerializer(serializers.ModelSerializer):

    class Meta:
        model = OrderRecipe
        fields = (
            'id',
            'recipe_id',
            'order_id',
            'quantity',

        )

class OrderSerializer(serializers.ModelSerializer):
# recipe_list = RecipeSerializer(source='recipe', read_only=True, many=True)
    created = serializers.DateTimeField(format='%Y-%m-%d %H:%M:%S')
    class Meta:
        model = Order
        fields = (
            'id',
            'name',
            'table',
            'price',
            'status',
            'slug',
            'deleted',
            'created',
        )


import torch
from torch import nn
from models.sh import eval_sh

class Embedding(nn.Module):
    def __init__(self, in_channels, N_freqs, logscale=True):
        """
        Defines a function that embeds x to (x, sin(2^k x), cos(2^k x), ...)
        in_channels: number of input channels (3 for both xyz and direction)
        """
        super(Embedding, self).__init__()
        self.N_freqs = N_freqs
        self.in_channels = in_channels
        self.funcs = [torch.sin, torch.cos]
        self.out_channels = in_channels*(len(self.funcs)*N_freqs+1)

        if logscale:
            self.freq_bands = 2**torch.linspace(0, N_freqs-1, N_freqs)
        else:
            self.freq_bands = torch.linspace(1, 2**(N_freqs-1), N_freqs)

    def forward(self, x):
        """
        Embeds x to (x, sin(2^k x), cos(2^k x), ...)
        Different from the paper, "x" is also in the output
        See https://github.com/bmild/nerf/issues/12

        Inputs:
            x: (B, self.in_channels)

        Outputs:
            out: (B, self.out_channels)
        """
        out = [x]
        for freq in self.freq_bands:
            for func in self.funcs:
                out += [func(freq*x)]

        return torch.cat(out, -1)

class NeRF(nn.Module):
    def __init__(self,
                 D=8, W=256,
                 in_channels_xyz=63, in_channels_dir=27,
                 skips=[4], deg=2):
        """
        D: number of layers for density (sigma) encoder
        W: number of hidden units in each layer
        in_channels_xyz: number of input channels for xyz (3+3*10*2=63 by default)
        in_channels_dir: number of input channels for direction (3+3*4*2=27 by default)
        skips: add skip connection in the Dth layer
        """
        super(NeRF, self).__init__()
        self.D = D
        self.W = W
        self.in_channels_xyz = in_channels_xyz
        self.in_channels_dir = in_channels_dir
        self.skips = skips
        self.deg = deg

        # xyz encoding layers
        for i in range(D):
            if i == 0:
                layer = nn.Linear(in_channels_xyz, W)
            elif i in skips:
                layer = nn.Linear(W+in_channels_xyz, W)
            else:
                layer = nn.Linear(W, W)
            layer = nn.Sequential(layer, nn.ReLU(True))
            setattr(self, f"xyz_encoding_{i+1}", layer)
        # self.xyz_encoding_final = nn.Linear(W, W)

        # # direction encoding layers
        # self.dir_encoding = nn.Sequential(
        #                         nn.Linear(W+in_channels_dir, W),
        #                         nn.ReLU(True))

        # output layers
        self.sigma = nn.Sequential(nn.Linear(W, W),
                                   nn.ReLU(True),
                                   nn.Linear(W, 1))
        # self.sh = nn.Linear(W, 3 * (self.deg + 1)**2)
        self.sh = nn.Sequential(nn.Linear(W, W),
                                   nn.ReLU(True),
                                   nn.Linear(W, 3 * (self.deg + 1)**2))

    def forward(self, x, dirs=None, sigma_sh_only=False):
        """
        Encodes input (xyz+dir) to sh+sigma (not ready to render yet).
        For rendering this ray, please see rendering.py

        Inputs:
            x: (B, self.in_channels_xyz(+self.in_channels_dir))
               the embedded vector of position and direction
            sigma_only: whether to infer sigma only. If True,
                        x is of shape (B, self.in_channels_xyz)

        Outputs:
            if sigma_ony:
                sigma: (B, 1) sigma
            else:
                out: (B, 4), sh and sigma
        """
        input_xyz = x

        xyz_ = input_xyz
        for i in range(self.D):
            if i in self.skips:
                xyz_ = torch.cat([input_xyz, xyz_], -1)
            xyz_ = getattr(self, f"xyz_encoding_{i+1}")(xyz_)

        sigma = self.sigma(xyz_)
        sh = self.sh(xyz_)

        if sigma_sh_only:
            out = torch.cat([sigma, sh], -1)
            return out

        rgb = eval_sh(deg=self.deg, sh=sh.reshape(-1, 3, (self.deg + 1)**2), dirs=dirs) # sh: [..., C, (deg + 1) ** 2]
        rgb = torch.sigmoid(rgb)

        # if extract_time:
        out = torch.cat([sigma, rgb, sh], -1)
        return out


import numpy as np
from scipy.spatial.transform import Rotation

class SE3Control(object):
    """

    """
    def __init__(self, quad_params):
        """
        Parameters:
            quad_params, dict with keys specified in rotorpy/vehicles
        """

        # Quadrotor physical parameters.
        # Inertial parameters
        self.mass            = quad_params['mass'] # kg
        self.Ixx             = quad_params['Ixx']  # kg*m^2
        self.Iyy             = quad_params['Iyy']  # kg*m^2
        self.Izz             = quad_params['Izz']  # kg*m^2
        self.Ixy             = quad_params['Ixy']  # kg*m^2
        self.Ixz             = quad_params['Ixz']  # kg*m^2
        self.Iyz             = quad_params['Iyz']  # kg*m^2

        # Frame parameters
        self.c_Dx            = quad_params['c_Dx']  # drag coeff, N/(m/s)**2
        self.c_Dy            = quad_params['c_Dy']  # drag coeff, N/(m/s)**2
        self.c_Dz            = quad_params['c_Dz']  # drag coeff, N/(m/s)**2

        self.num_rotors      = quad_params['num_rotors']
        self.rotor_pos       = quad_params['rotor_pos']

        # Rotor parameters
        self.rotor_speed_min = quad_params['rotor_speed_min'] # rad/s
        self.rotor_speed_max = quad_params['rotor_speed_max'] # rad/s

        self.k_eta           = quad_params['k_eta']     # thrust coeff, N/(rad/s)**2
        self.k_m             = quad_params['k_m']       # yaw moment coeff, Nm/(rad/s)**2
        self.k_d             = quad_params['k_d']       # rotor drag coeff, N/(m/s)
        self.k_z             = quad_params['k_z']       # induced inflow coeff N/(m/s)
        self.k_flap          = quad_params['k_flap']    # Flapping moment coefficient Nm/(m/s)

        # Motor parameters
        self.tau_m           = quad_params['tau_m']     # motor reponse time, seconds

        # You may define any additional constants you like including control gains.
        self.inertia = np.array([[self.Ixx, self.Ixy, self.Ixz],
                                 [self.Ixy, self.Iyy, self.Iyz],
                                 [self.Ixz, self.Iyz, self.Izz]]) # kg*m^2
        self.g = 9.81 # m/s^2

        # Gains
        self.kp_pos = np.array([6.5,6.5,15])
        self.kd_pos = np.array([4.0, 4.0, 9])
        self.kp_att = 544
        self.kd_att = 46.64

        # Linear map from individual rotor forces to scalar thrust and vector
        # moment applied to the vehicle.
        k = self.k_m/self.k_eta

        # Below is an automated generation of the control allocator matrix. It assumes that all thrust vectors are aligned
        # with the z axis and that the "sign" of each rotor yaw moment alternates starting with positive for r1.
        self.f_to_TM = np.vstack((np.ones((1,self.num_rotors)),np.hstack([np.cross(self.rotor_pos[key],np.array([0,0,1])).reshape(-1,1)[0:2] for key in self.rotor_pos]), np.array([k*(-1)**i for i in range(self.num_rotors)]).reshape(1,-1)))
        self.TM_to_f = np.linalg.inv(self.f_to_TM)

    def update_ref(self, t, flat_output):
        """
        This function receives the current time, and desired flat
        outputs. It returns the reference command inputs.
        Follows https://repository.upenn.edu/edissertations/547/

        Inputs:
            t, present time in seconds
            flat_output, a dict describing the present desired flat outputs with keys
                x,        position, m
                x_dot,    velocity, m/s
                x_ddot,   acceleration, m/s**2  a
                x_dddot,  jerk, m/s**3          a_dot
                x_ddddot, snap, m/s**4          a_ddot
                yaw,      yaw angle, rad
                yaw_dot,  yaw rate, rad/s
                yaw_ddot, yaw acceleration, rad/s**2  #required! not the same if computing command using controller

        Outputs:
            control_input, a dict describing the present computed control inputs with keys
                cmd_motor_speeds, rad/s
                cmd_thrust, N (for debugging and laboratory; not used by simulator)
                cmd_moment, N*m (for debugging; not used by simulator)
                cmd_q, quaternion [i,j,k,w] (for laboratory; not used by simulator)
                cmd_w, angular velocity
                cmd_a, angular acceleration
        """
        cmd_motor_speeds = np.zeros((4,))
        cmd_q = np.zeros((4,))

        def normalize(x):
            """Return normalized vector."""
            return x / np.linalg.norm(x)

        # def vee_map(S):
        #     """Return vector corresponding to given skew symmetric matrix."""
        #     return np.array([-S[1,2], S[0,2], -S[0,1]])

        # Desired force vector.
        t = flat_output['x_ddot']+ np.array([0, 0, self.g])
        b3 = normalize(t)
        F_des = self.mass * (t)# this is vectorized

        # Desired thrust is force projects onto b3 axis.
        # R = Rotation.from_quat(state['q']).as_matrix() #this is where most of the problem is, there is no error in rotation!
        # b3 = R @ np.array([0, 0, 1])
        u1 = np.dot(F_des, b3)

        # Desired orientation to obtain force vector.
        b3_des = normalize(F_des) #b3_des and b3 are the same
        yaw_des = flat_output['yaw']
        c1_des = np.array([np.cos(yaw_des), np.sin(yaw_des), 0])
        b2_des = normalize(np.cross(b3_des, c1_des))
        b1_des = np.cross(b2_des, b3_des)
        R_des = np.stack([b1_des, b2_des, b3_des]).T

        R = R_des# assume we have perfect tracking on rotation
        # Orientation error.
        # S_err = 0.5 * (R_des.T @ R - R.T @ R_des)
        # att_err = vee_map(S_err)

        # Following section follows Mellinger paper to compute reference angular velocity
        dot_u1 = np.dot(b3,flat_output['x_dddot'])
        hw = self.mass/u1*(flat_output['x_dddot']-dot_u1*b3)
        p  = np.dot(-hw, b2_des)
        q  = np.dot(hw, b1_des)
        w_des = np.array([0, 0, flat_output['yaw_dot']])
        r  = np.dot(w_des, b3_des)
        Omega = np.array([p, q, r])

        wwu1b3 = np.cross(Omega, np.cross(Omega, u1*b3))
        ddot_u1 = np.dot(b3, self.mass*flat_output['x_ddddot']) - np.dot(b3, wwu1b3)
        ha = 1.0/u1*(self.mass*flat_output['x_ddddot'] - ddot_u1*b3 - 2*np.cross(Omega,dot_u1*b3) - wwu1b3)
        p_dot = np.dot(-ha, b2_des)
        q_dot = np.dot(ha, b1_des)
        np.cross(Omega, Omega)
        r_dot = flat_output['yaw_ddot'] *np.dot(np.array([0,0,1.0]), b3_des) #uniquely need yaw_ddot
        Alpha = np.array([p_dot, q_dot, r_dot])

        u2 =  self.inertia @ Alpha + np.cross(Omega, self.inertia @ Omega)
        # print(u1,u2)
        TM = np.array([u1, u2[0], u2[1], u2[2]])
        cmd_motor_forces = self.TM_to_f @ TM
        cmd_motor_speeds = cmd_motor_forces / self.k_eta
        cmd_motor_speeds = np.sign(cmd_motor_speeds) * np.sqrt(np.abs(cmd_motor_speeds))

        cmd_q = Rotation.from_matrix(R_des).as_quat()

        control_input = {'cmd_motor_speeds':cmd_motor_speeds,
                        'cmd_thrust':u1,
                        'cmd_moment':u2,
                        'cmd_q':cmd_q,
                        'cmd_w':Omega,
                        'cmd_a':Alpha}
        return control_input

    def update(self, t, state, flat_output):
        """
        This function receives the current time, true state, and desired flat
        outputs. It returns the command inputs.

        Inputs:
            t, present time in seconds
            state, a dict describing the present state with keys
                x, position, m
                v, linear velocity, m/s
                q, quaternion [i,j,k,w]
                w, angular velocity, rad/s
            flat_output, a dict describing the present desired flat outputs with keys
                x,        position, m
                x_dot,    velocity, m/s
                x_ddot,   acceleration, m/s**2
                x_dddot,  jerk, m/s**3
                x_ddddot, snap, m/s**4
                yaw,      yaw angle, rad
                yaw_dot,  yaw rate, rad/s

        Outputs:
            control_input, a dict describing the present computed control inputs with keys
                cmd_motor_speeds, rad/s
                cmd_thrust, N
                cmd_moment, N*m
                cmd_q, quaternion [i,j,k,w]
        """
        cmd_motor_speeds = np.zeros((4,))
        cmd_thrust = 0
        cmd_moment = np.zeros((3,))
        cmd_q = np.zeros((4,))

        def normalize(x):
            """Return normalized vector."""
            return x / np.linalg.norm(x)

        def vee_map(S):
            """Return vector corresponding to given skew symmetric matrix."""
            return np.array([-S[1,2], S[0,2], -S[0,1]])

        # Desired force vector.
        pos_err  = state['x'] - flat_output['x']
        dpos_err = state['v'] - flat_output['x_dot']
        F_des = self.mass * (- self.kp_pos*pos_err
                             - self.kd_pos*dpos_err
                             + flat_output['x_ddot']
                             + np.array([0, 0, self.g]))

        # Desired thrust is force projects onto b3 axis.
        R = Rotation.from_quat(state['q']).as_matrix()
        b3 = R @ np.array([0, 0, 1])
        u1 = np.dot(F_des, b3)

        # Desired orientation to obtain force vector.
        b3_des = normalize(F_des)
        yaw_des = flat_output['yaw']
        c1_des = np.array([np.cos(yaw_des), np.sin(yaw_des), 0])
        b2_des = normalize(np.cross(b3_des, c1_des))
        b1_des = np.cross(b2_des, b3_des)
        R_des = np.stack([b1_des, b2_des, b3_des]).T

        # Orientation error.
        S_err = 0.5 * (R_des.T @ R - R.T @ R_des)
        att_err = vee_map(S_err)

        # Angular velocity error (this is oversimplified).
        w_des = np.array([0, 0, flat_output['yaw_dot']])
        w_err = state['w'] - w_des

        # Angular control; vector units of N*m.
        u2 = self.inertia @ (-self.kp_att*att_err - self.kd_att*w_err)

        # Compute motor speeds. Avoid taking square root of negative numbers.
        TM = np.array([u1, u2[0], u2[1], u2[2]])
        cmd_motor_forces = self.TM_to_f @ TM
        cmd_motor_speeds = cmd_motor_forces / self.k_eta
        cmd_motor_speeds = np.sign(cmd_motor_speeds) * np.sqrt(np.abs(cmd_motor_speeds))

        cmd_thrust = u1
        cmd_moment = u2
        cmd_q = Rotation.from_matrix(R_des).as_quat()

        control_input = {'cmd_motor_speeds':cmd_motor_speeds,
                         'cmd_thrust':cmd_thrust,
                         'cmd_moment':cmd_moment,
                         'cmd_q':cmd_q}
        return control_input


import re
from ipykernel.ipkernel import IPythonKernel
from .shell import OverrideShell

class NowKernel(IPythonKernel):

    def __init__(self, **kwargs):
        super(NowKernel, self).__init__(**kwargs)
        OverrideShell(self.shell)

    def do_execute(self, code, silent, store_history=True,
                   user_expressions=None, allow_stdin=False):
        reply_content = super(NowKernel, self).do_execute(
            code, silent, store_history=store_history,
            user_expressions=user_expressions, allow_stdin=allow_stdin)
        return reply_content


# -*- coding: utf-8 -*-
"""
Created on Mon Oct 28 08:19:12 2019

@author: rastf
"""

# -*- coding: utf-8 -*-
"""
Created on Mon Oct 21 07:26:54 2019

@author: rastf
"""

from astropy.io import fits
from astropy.wcs import WCS
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.gridspec as gridspec
from matplotlib.colors import LogNorm
from matplotlib.patches import Circle
from astropy.visualization import simple_norm

#finding out what the coordinate system for each image is
wcs24 = WCS(fits.getheader('G54_Spitzer_24.fits'))
wcs70 = WCS(fits.getheader('G54_PACS_70.fits'))
wcs160 = WCS(fits.getheader('G54_PACS_160.fits'))
wcs250 = WCS(fits.getheader('G54_SPIRE_250.fits'))
wcs350 = WCS(fits.getheader('G54_SPIRE_350.fits'))
wcs500 = WCS(fits.getheader('G54_SPIRE_500.fits'))
wcs870 = WCS(fits.getheader('G54.1_FINAL_870.fits'))

#setting up the grid for plotting
f= plt.figure(figsize = (5,5))
gs1 = gridspec.GridSpec(1,1)
gs1.update(wspace=0.025, hspace=0)

#importing the data in
SPITZER24=fits.getdata('G54_Spitzer_24.fits')
PACS70=fits.getdata('G54_PACS_70.fits')
PACS160=fits.getdata('G54_PACS_160.fits')
SPIRE250=fits.getdata('G54_SPIRE_250.fits')
SPIRE350=fits.getdata('G54_SPIRE_350.fits')
SPIRE500=fits.getdata('G54_SPIRE_500.fits')
FINAL870=fits.getdata('G54.1_FINAL_870.fits')

#setting vmin and vmax
vmin=-0.03
vmax=1
cmap='cubehelix'
#defining some lists that will be needed for plotting
ax = []
lat = []
lon = []

# =============================================================================
# SPITZER 24 IMAGE
# =============================================================================
t=0
ax.append(plt.subplot(gs1[t], projection=wcs24)) #creating subplot and assigning correct wcs
img = ax[t].imshow(SPITZER24, cmap=cmap, origin = 'lower', norm=simple_norm(SPITZER24)) #plotting image with LogNorm
ax[t].text(0.5,0.92,'Spitzer MIPS 24$\mu$m', fontsize=13, family='serif', ha = 'left', va = 'center', transform = ax[t].transAxes, color='white') #adding text
c = Circle((292.6167, 18.8683), 0.025, edgecolor='white', facecolor='none',alpha=0.5,transform=ax[t].get_transform('fk5')) #defining circle
p = Circle((292.6292, 18.8667), 0.001, edgecolor='white', facecolor='none',alpha=0.5,transform=ax[t].get_transform('fk5')) #defining circle
ax[t].add_patch(c) #adding the circle
ax[t].add_patch(p)
#c2 = Circle((292.6167, 18.8683), 0.022, edgecolor='white', facecolor='none',alpha=0.5,transform=ax[t].get_transform('fk5')) #defining circle
#ax[t].add_patch(c2) #adding the circle2
lon.append(ax[t].coords[0]) #defining x axis
lat.append(ax[t].coords[1]) #defining y axis
lon[t].set_ticklabel_visible(True) #choosing to show or hide the x axis
lon[t].set_ticks_visible(True)
lat[t].set_ticklabel_visible(True) #choosing to show or hide the y axis
lat[t].set_ticks_visible(True)
ax[t].set_xlim((51,101)) #restricting field of view
ax[t].set_ylim((51,101)) #restricting field of view


import os
import logging
import discord
from discord.ext import commands
from dotenv import load_dotenv
import utils

load_dotenv()

logger = logging.getLogger('discord')
logger.setLevel(logging.INFO)
handler = logging.FileHandler(filename='discord.log', encoding='utf-8', mode='w')
handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
logger.addHandler(handler)

client = discord.Client()

@client.event
async def on_ready():
    logger.info(f"logged in as {client.user}")

@client.event
async def on_message(ctx):
    content = ctx.content
    if ctx.author == client.user:
        logger.info("message sent by bot itself, ignoring...")
        return

    if content.startswith('!o'):
        fields = content.split(' ')
        try:
            dice_pool = int(fields[1])
            modifier = int(fields[2])
            dt = int(fields[3])
        except IndexError as error:
            logger.warning(f"invalid fields for {ctx.author} message threw the exception bellow...\n {error}")
            await ctx.channel.send("Campos inválidos!")
            return

        dice_info = utils.dice_roll(dice_pool,modifier,dt)

        embed_image = discord.Embed(
            title = f'{ctx.author} rolou...',
            color=dice_info[3],
        )
        embed_image.add_field(name="**Dados**", value=dice_info[0], inline=True)
        embed_image.add_field(name="**Modificador**", value=modifier, inline=True)
        embed_image.add_field(name="**Dificuldade**", value=dt, inline=True)
        embed_image.add_field(name="**Rolagem final:**", value=dice_info[5]+modifier)
        embed_image.add_field(name="**Resultado:**", value=dice_info[4])
        embed_image.set_thumbnail(url=dice_info[2])

        logger.info(f"rolling dice for {ctx.author}...")
        await ctx.channel.send(embed=embed_image)
        return

    elif content.startswith('ordo!ping'):
        embed_image = discord.Embed(
            title = 'Ping...',
            color=0x492ea4,
        )
        logger.info(f'running ping command by {ctx.author}...')
        embed_image.add_field(name="**Pong!**", value=f'{round(client.latency * 1000)}ms')
        await ctx.channel.send(embed=embed_image)
        return

    elif content.startswith('!ro'):
        fields=content.split(' ')
        dice_info = utils.other_dice(fields[1])
        embed_image = discord.Embed(
            color=0x492ea4
        )
        embed_image.add_field(name="**Dados**", value=dice_info[0], inline=True)
        embed_image.add_field(name="**Total**", value=dice_info[3], inline=True)
        await ctx.channel.send(embed=embed_image)
        return

if __name__ == "__main__":
    client.run(os.environ["BOT_TOKEN"])


from typing import Tuple

import torch

from .base_trainer import BaseTrainer
from .utils import get_batch_tensors, get_dict

class StandardTrainer(BaseTrainer):

    def __init__(self, model, optimizer, loss, metrics=None,
                 x_device=None, x_type=torch.float, y_device=None, y_type=torch.long):
        """
        :param model: torch module
        :param optimizer: torch optimizer
        :param loss: loss function with signature function(preds, trues)
        :param metrics: list of metric functions with signature function(preds, trues)
        :param x_device: device to put inputs
        :param x_type: type to cast inputs
        :param y_device: device to put labels
        :param y_type: type to cast labels
        """
        super().__init__([model], [optimizer], loss, metrics)

        self.x_device = x_device
        self.x_type = x_type
        self.y_device = y_device
        self.y_type = y_type

        self.training = True

    def _extract_data(self, batch_data):
        return get_batch_tensors(batch_data, self.x_type, self.x_device, self.y_type, self.y_device)

    def _train_one_batch(self, inputs, targets) -> Tuple[dict, object]:
        if isinstance(inputs, dict):
            preds = self.models[0](**inputs)
        elif isinstance(inputs, list):
            preds = self.models[0](*inputs)
        else:
            preds = self.models[0](inputs)

        loss = self.loss(preds, targets)

        if isinstance(loss, dict):
            loss_dict = {k: float(loss[k]) for k in loss}
            loss = loss["Loss"]
        else:
            loss_dict = {"Loss": float(loss)}

        loss.backward()
        self.optimizers[0].step()

        return loss_dict, preds

    def _get_train_measures(self, inputs, targets, loss_dict, cache) -> dict:
        preds = cache

        measures = loss_dict
        for m in self.metrics:
            m_value = m(preds, targets)
            m_dict = get_dict(m_value, name=m._get_name())
            measures.update(m_dict)

        return measures

    def _get_eval_cache(self, inputs, targets):
        if isinstance(inputs, dict):
            preds = self.models[0](**inputs)
        elif isinstance(inputs, list):
            preds = self.models[0](*inputs)
        else:
            preds = self.models[0](inputs)

        return preds, targets

    def _get_eval_logs(self, eval_caches):
        preds, trues = eval_caches

        loss = self.loss(preds, trues)
        loss_dict = get_dict(loss, prefix="Val_", name="Loss")

        measures = loss_dict
        for m in self.metrics:
            m_value = m(preds, trues)
            m_dict = get_dict(m_value, prefix="val ", name=m._get_name())
            measures.update(m_dict)

        return measures

    def extra_repr(self):
        return f"x_device={self.x_device}, x_type={self.x_type}, " \
               f"y_device={self.y_device}, y_type={self.y_type}"


# https://github.com/skyhehe123/VoxelNet-pytorch

# voxelnet

import torch.utils.data as data
import time

import torch.optim as optim
import torch.nn.init as init
import torch.backends.cudnn
from __future__ import division
import numpy as np
import math
import mayavi.mlab as mlab
import cv2
import cv2

from torch.autograd import Variable
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

import os
import cv2
import matplotlib.pyplot as plt
import torch.utils.data as data
import torch.nn.functional as F
import numpy as np
import torch.backends.cudnn
import cv2
import matplotlib.pyplot as plt
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.enabled = True

# voxel configuration

class config:

    # classes
    class_list = ['Car', 'Bicycle', 'Pedestrian', 'Truck',
                  'Small vehicles', 'Traffic signal', 'Traffic sign',
                  'Utility vehicle', 'Sidebars', 'Speed bumper',
                  'Curbstone', 'Solid line', 'Irrelevant signs',
                  'Road blocks', 'Tractor', 'Non-drivable street',
                  'Zebra crossing', 'Obstacles / trash', 'Poles',
                  'RD restricted area', 'Animals', 'Grid structure',
                  'Signal corpus', 'Drivable cobblestone', 'Electronic traffic',
                  'Slow drive area', 'Nature object', 'Parking area',
                  'Sidewalk', 'Ego car', 'Painted driv. instr.',
                  'Traffic guide obj.', 'Dashed line', 'RD normal street',
                  'Sky', 'Buildings', 'Blurred area', 'Rain dirt']

    # batch size
    N = 2

    # maxiumum number of points per voxel
    T = 35

    # voxel size
    vd = 0.4
    vh = 0.2
    vw = 0.2

    # points cloud range
    xrange = (0, 100)
    yrange = (-50, 50)
    zrange = (-5, 5)

    # voxel grid
    W = math.ceil((xrange[1] - xrange[0]) / vw)
    H = math.ceil((yrange[1] - yrange[0]) / vh)
    D = math.ceil((zrange[1] - zrange[0]) / vd)

    # iou threshold
    pos_threshold = 0.6
    neg_threshold = 0.45

    #   anchors: (200, 176, 2, 7) x y z h w l r
    x = np.linspace(xrange[0]+vw, xrange[1]-vw, W/2)
    y = np.linspace(yrange[0]+vh, yrange[1]-vh, H/2)
    cx, cy = np.meshgrid(x, y)
    # all is (w, l, 2)
    cx = np.tile(cx[..., np.newaxis], 2)
    cy = np.tile(cy[..., np.newaxis], 2)
    cz = np.ones_like(cx) * -1.0
    w = np.ones_like(cx) * 1.6
    l = np.ones_like(cx) * 3.9
    h = np.ones_like(cx) * 1.56
    r = np.ones_like(cx)
    r[..., 0] = 0
    r[..., 1] = np.pi/2
    anchors = np.stack([cx, cy, cz, h, w, l, r], axis=-1)

    anchors_per_position = 2

    # non-maximum suppression
    nms_threshold = 0.1
    score_threshold = 0.96

# _____________________________________________________________________________________________________ #

class A2D2Dataset(data.Dataset):

    def __init__(self, root='./KITTI',set='train',type='velodyne_train'):
        self.type = type
        self.root = root
        self.data_path = os.path.join(root, 'training')
        self.lidar_path = os.path.join(self.data_path, "crop/")
        self.image_path = os.path.join(self.data_path, "image_2/")
        self.calib_path = os.path.join(self.data_path, "calib/")
        self.label_path = os.path.join(self.data_path, "label_2/")

        with open(os.path.join(self.data_path, '%s.txt' % set)) as f:
            self.file_list = f.read().splitlines()

        self.T = config.T
        self.vd = config.vd
        self.vh = config.vh
        self.vw = config.vw
        self.xrange = config.xrange
        self.yrange = config.yrange
        self.zrange = config.zrange
        self.anchors = config.anchors.reshape(-1,7)
        self.feature_map_shape = (int(config.H / 2), int(config.W / 2))
        self.anchors_per_position = config.anchors_per_position
        self.pos_threshold = config.pos_threshold
        self.neg_threshold = config.neg_threshold

    def cal_target(self, gt_box3d):
        # Input:
        #   labels: (N,)
        #   feature_map_shape: (w, l)
        #   anchors: (w, l, 2, 7)
        # Output:
        #   pos_equal_one (w, l, 2)
        #   neg_equal_one (w, l, 2)
        #   targets (w, l, 14)
        # attention: cal IoU on birdview

        anchors_d = np.sqrt(self.anchors[:, 4] ** 2 + self.anchors[:, 5] ** 2)

        pos_equal_one = np.zeros((*self.feature_map_shape, 2))
        neg_equal_one = np.zeros((*self.feature_map_shape, 2))
        targets = np.zeros((*self.feature_map_shape, 14))

        gt_xyzhwlr = box3d_corner_to_center_batch(gt_box3d)

        anchors_corner = anchors_center_to_corner(self.anchors)

        anchors_standup_2d = corner_to_standup_box2d_batch(anchors_corner)
        # BOTTLENECK
        gt_standup_2d = corner_to_standup_box2d_batch(gt_box3d)

        iou = bbox_overlaps(
            np.ascontiguousarray(anchors_standup_2d).astype(np.float32),
            np.ascontiguousarray(gt_standup_2d).astype(np.float32),
        )

        id_highest = np.argmax(iou.T, axis=1)  # the maximum anchor's ID
        id_highest_gt = np.arange(iou.T.shape[0])
        mask = iou.T[id_highest_gt, id_highest] > 0
        id_highest, id_highest_gt = id_highest[mask], id_highest_gt[mask]
        # find anchor iou > cfg.XXX_POS_IOU
        id_pos, id_pos_gt = np.where(iou > self.pos_threshold)
        # find anchor iou < cfg.XXX_NEG_IOU
        id_neg = np.where(np.sum(iou < self.neg_threshold,
                                 axis=1) == iou.shape[1])[0]

        id_pos = np.concatenate([id_pos, id_highest])
        id_pos_gt = np.concatenate([id_pos_gt, id_highest_gt])
        # TODO: uniquify the array in a more scientific way
        id_pos, index = np.unique(id_pos, return_index=True)
        id_pos_gt = id_pos_gt[index]
        id_neg.sort()
        # cal the target and set the equal one
        index_x, index_y, index_z = np.unravel_index(
            id_pos, (*self.feature_map_shape, self.anchors_per_position))
        pos_equal_one[index_x, index_y, index_z] = 1
        # ATTENTION: index_z should be np.array

        targets[index_x, index_y, np.array(index_z) * 7] = \
            (gt_xyzhwlr[id_pos_gt, 0] - self.anchors[id_pos, 0]) / anchors_d[id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 1] = \
            (gt_xyzhwlr[id_pos_gt, 1] - self.anchors[id_pos, 1]) / anchors_d[id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 2] = \
            (gt_xyzhwlr[id_pos_gt, 2] - self.anchors[id_pos, 2]) / self.anchors[id_pos, 3]
        targets[index_x, index_y, np.array(index_z) * 7 + 3] = np.log(
            gt_xyzhwlr[id_pos_gt, 3] / self.anchors[id_pos, 3])
        targets[index_x, index_y, np.array(index_z) * 7 + 4] = np.log(
            gt_xyzhwlr[id_pos_gt, 4] / self.anchors[id_pos, 4])
        targets[index_x, index_y, np.array(index_z) * 7 + 5] = np.log(
            gt_xyzhwlr[id_pos_gt, 5] / self.anchors[id_pos, 5])
        targets[index_x, index_y, np.array(index_z) * 7 + 6] = (
                gt_xyzhwlr[id_pos_gt, 6] - self.anchors[id_pos, 6])
        index_x, index_y, index_z = np.unravel_index(
            id_neg, (*self.feature_map_shape, self.anchors_per_position))
        neg_equal_one[index_x, index_y, index_z] = 1
        # to avoid a box be pos/neg in the same time
        index_x, index_y, index_z = np.unravel_index(
            id_highest, (*self.feature_map_shape, self.anchors_per_position))
        neg_equal_one[index_x, index_y, index_z] = 0

        return pos_equal_one, neg_equal_one, targets

    def preprocess(self, lidar):

        # shuffling the points
        np.random.shuffle(lidar)

        voxel_coords = ((lidar[:, :3] - np.array([self.xrange[0], self.yrange[0], self.zrange[0]])) / (
                        self.vw, self.vh, self.vd)).astype(np.int32)

        # convert to  (D, H, W)
        voxel_coords = voxel_coords[:,[2,1,0]]
        voxel_coords, inv_ind, voxel_counts = np.unique(voxel_coords, axis=0, \
                                                  return_inverse=True, return_counts=True)

        voxel_features = []

        for i in range(len(voxel_coords)):
            voxel = np.zeros((self.T, 7), dtype=np.float32)
            pts = lidar[inv_ind == i]
            if voxel_counts[i] > self.T:
                pts = pts[:self.T, :]
                voxel_counts[i] = self.T
            # augment the points
            voxel[:pts.shape[0], :] = np.concatenate((pts, pts[:, :3] - np.mean(pts[:, :3], 0)), axis=1)
            voxel_features.append(voxel)
        return np.array(voxel_features), voxel_coords

    def __getitem__(self, i):

        lidar_file = self.lidar_path + '/' + self.file_list[i] + '.bin'
        calib_file = self.calib_path + '/' + self.file_list[i] + '.txt'
        label_file = self.label_path + '/' + self.file_list[i] + '.txt'
        image_file = self.image_path + '/' + self.file_list[i] + '.png'

        calib = load_kitti_calib(calib_file)
        Tr = calib['Tr_velo2cam']
        gt_box3d = load_kitti_label(label_file, Tr)
        lidar = np.fromfile(lidar_file, dtype=np.float32).reshape(-1, 4)

        if self.type == 'velodyne_train':
            image = cv2.imread(image_file)

            # data augmentation
            lidar, gt_box3d = aug_data(lidar, gt_box3d)

            # specify a range
            lidar, gt_box3d = get_filtered_lidar(lidar, gt_box3d)

            # voxelize
            voxel_features, voxel_coords = self.preprocess(lidar)

            # bounding-box encoding
            pos_equal_one, neg_equal_one, targets = self.cal_target(gt_box3d)

            return voxel_features, voxel_coords, pos_equal_one, neg_equal_one, targets, image, calib, self.file_list[i]

        elif self.type == 'velodyne_test':
            NotImplemented

        else:
            raise ValueError('the type invalid')

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

# _____________________________________________________________________________________________________ #

# conv2d + bn + relu
class Conv2d(nn.Module):

    def __init__(self, in_channels, out_channels, k, s, p, activation=True, batch_norm=True):
        super(Conv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=k, stride=s, padding=p)
        if batch_norm:
            self.bn = nn.BatchNorm2d(out_channels)
        else:
            self.bn = None
        self.activation = activation

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x=self.bn(x)
        if self.activation:
            return F.relu(x, inplace=True)
        else:
            return x

# conv3d + bn + relu
class Conv3d(nn.Module):

    def __init__(self, in_channels, out_channels, k, s, p, batch_norm=True):
        super(Conv3d, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=k, stride=s, padding=p)
        if batch_norm:
            self.bn = nn.BatchNorm3d(out_channels)
        else:
            self.bn = None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)

        return F.relu(x, inplace=True)

# Fully Connected Network
class FCN(nn.Module):

    def __init__(self, cin, cout):
        super(FCN, self).__init__()
        self.cout = cout
        self.linear = nn.Linear(cin, cout)
        self.bn = nn.BatchNorm1d(cout)

    def forward(self, x):
        # KK is the stacked k across batch
        kk, t, _ = x.shape
        x = self.linear(x.view(kk*t, -1))
        x = F.relu(self.bn(x))
        return x.view(kk, t, -1)

# Voxel Feature Encoding layer
class VFE(nn.Module):

    def __init__(self,cin,cout):
        super(VFE, self).__init__()
        assert cout % 2 == 0
        self.units = cout // 2
        self.fcn = FCN(cin,self.units)

    def forward(self, x, mask):
        # point-wise feature
        pwf = self.fcn(x)
        # locally aggregated feature
        laf = torch.max(pwf, 1)[0].unsqueeze(1).repeat(1, config.T, 1)
        # point-wise concat feature
        pwcf = torch.cat((pwf, laf), dim=2)
        # apply mask
        mask = mask.unsqueeze(2).repeat(1, 1, self.units * 2)
        pwcf = pwcf * mask.float()

        return pwcf

# Stacked Voxel Feature Encoding
class SVFE(nn.Module):

    def __init__(self):
        super(SVFE, self).__init__()
        self.vfe_1 = VFE(7, 32)
        self.vfe_2 = VFE(32, 128)
        self.fcn = FCN(128, 128)

    def forward(self, x):
        mask = torch.ne(torch.max(x, 2)[0], 0)
        x = self.vfe_1(x, mask)
        x = self.vfe_2(x, mask)
        x = self.fcn(x)
        # element-wise max pooling
        x = torch.max(x, 1)[0]
        return x

# Convolutional Middle Layer
class CML(nn.Module):
    def __init__(self):
        super(CML, self).__init__()
        self.conv3d_1 = Conv3d(128, 64, 3, s=(2, 1, 1), p=(1, 1, 1))
        self.conv3d_2 = Conv3d(64, 64, 3, s=(1, 1, 1), p=(0, 1, 1))
        self.conv3d_3 = Conv3d(64, 64, 3, s=(2, 1, 1), p=(1, 1, 1))

    def forward(self, x):
        x = self.conv3d_1(x)
        x = self.conv3d_2(x)
        x = self.conv3d_3(x)
        return x

# Region Proposal Network
class RPN(nn.Module):
    def __init__(self):
        super(RPN, self).__init__()
        self.block_1 = [Conv2d(128, 128, 3, 2, 1)]
        self.block_1 += [Conv2d(128, 128, 3, 1, 1) for _ in range(3)]
        self.block_1 = nn.Sequential(*self.block_1)

        self.block_2 = [Conv2d(128, 128, 3, 2, 1)]
        self.block_2 += [Conv2d(128, 128, 3, 1, 1) for _ in range(5)]
        self.block_2 = nn.Sequential(*self.block_2)

        self.block_3 = [Conv2d(128, 256, 3, 2, 1)]
        self.block_3 += [nn.Conv2d(256, 256, 3, 1, 1) for _ in range(5)]
        self.block_3 = nn.Sequential(*self.block_3)

        self.deconv_1 = nn.Sequential(nn.ConvTranspose2d(256, 256, 4, 4, 0), nn.BatchNorm2d(256))
        self.deconv_2 = nn.Sequential(nn.ConvTranspose2d(128, 256, 2, 2, 0), nn.BatchNorm2d(256))
        self.deconv_3 = nn.Sequential(nn.ConvTranspose2d(128, 256, 1, 1, 0), nn.BatchNorm2d(256))

        self.score_head = Conv2d(768, config.anchors_per_position, 1, 1, 0, activation=False, batch_norm=False)
        self.reg_head = Conv2d(768, 7 * config.anchors_per_position, 1, 1, 0, activation=False, batch_norm=False)

    def forward(self,x):
        x = self.block_1(x)
        x_skip_1 = x
        x = self.block_2(x)
        x_skip_2 = x
        x = self.block_3(x)
        x_0 = self.deconv_1(x)
        x_1 = self.deconv_2(x_skip_2)
        x_2 = self.deconv_3(x_skip_1)
        x = torch.cat((x_0,x_1,x_2),1)
        return self.score_head(x),self.reg_head(x)

class VoxelNet(nn.Module):

    def __init__(self):
        super(VoxelNet, self).__init__()
        self.svfe = SVFE()
        self.cml = CML()
        self.rpn = RPN()

    def voxel_indexing(self, sparse_features, coords):
        dim = sparse_features.shape[-1]

        dense_feature = Variable(torch.zeros(dim, config.N, config.D, config.H, config.W).cuda())

        dense_feature[:, coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]] = sparse_features

        return dense_feature.transpose(0, 1)

    def forward(self, voxel_features, voxel_coords):

        # feature learning network
        vwfs = self.svfe(voxel_features)
        vwfs = self.voxel_indexing(vwfs, voxel_coords)

        # convolutional middle network
        cml_out = self.cml(vwfs)

        # region proposal network

        # merge the depth and feature dim into one, output probability score map and regression map
        psm, rm = self.rpn(cml_out.view(config.N, -1, config.H, config.W))

        return psm, rm

# loss function

class VoxelLoss(nn.Module):
    def __init__(self, alpha, beta):
        super(VoxelLoss, self).__init__()
        self.smoothl1loss = nn.SmoothL1Loss(size_average=False)
        self.alpha = alpha
        self.beta = beta

    def forward(self, rm, psm, pos_equal_one, neg_equal_one, targets):

        p_pos = F.sigmoid(psm.permute(0,2,3,1))
        rm = rm.permute(0,2,3,1).contiguous()
        rm = rm.view(rm.size(0),rm.size(1),rm.size(2),-1,7)
        targets = targets.view(targets.size(0),targets.size(1),targets.size(2),-1,7)
        pos_equal_one_for_reg = pos_equal_one.unsqueeze(pos_equal_one.dim()).expand(-1,-1,-1,-1,7)

        rm_pos = rm * pos_equal_one_for_reg
        targets_pos = targets * pos_equal_one_for_reg

        cls_pos_loss = -pos_equal_one * torch.log(p_pos + 1e-6)
        cls_pos_loss = cls_pos_loss.sum() / (pos_equal_one.sum() + 1e-6)

        cls_neg_loss = -neg_equal_one * torch.log(1 - p_pos + 1e-6)
        cls_neg_loss = cls_neg_loss.sum() / (neg_equal_one.sum() + 1e-6)

        reg_loss = self.smoothl1loss(rm_pos, targets_pos)
        reg_loss = reg_loss / (pos_equal_one.sum() + 1e-6)
        conf_loss = self.alpha * cls_pos_loss + self.beta * cls_neg_loss
        return conf_loss, reg_loss

# _____________________________________________________________________________________________________ #

# _____________________________________________________________________________________________________ #

# data augmentation

def draw_polygon(img, box_corner, color = (255, 255, 255),thickness = 1):

    tup0 = (box_corner[0, 1],box_corner[0, 0])
    tup1 = (box_corner[1, 1],box_corner[1, 0])
    tup2 = (box_corner[2, 1],box_corner[2, 0])
    tup3 = (box_corner[3, 1],box_corner[3, 0])
    cv2.line(img, tup0, tup1, color, thickness, cv2.LINE_AA)
    cv2.line(img, tup1, tup2, color, thickness, cv2.LINE_AA)
    cv2.line(img, tup2, tup3, color, thickness, cv2.LINE_AA)
    cv2.line(img, tup3, tup0, color, thickness, cv2.LINE_AA)
    return img

def point_transform(points, tx, ty, tz, rx=0, ry=0, rz=0):
    # Input:
    #   points: (N, 3)
    #   rx/y/z: in radians
    # Output:
    #   points: (N, 3)
    N = points.shape[0]
    points = np.hstack([points, np.ones((N, 1))])
    mat1 = np.eye(4)
    mat1[3, 0:3] = tx, ty, tz
    points = np.matmul(points, mat1)
    if rx != 0:
        mat = np.zeros((4, 4))
        mat[0, 0] = 1
        mat[3, 3] = 1
        mat[1, 1] = np.cos(rx)
        mat[1, 2] = -np.sin(rx)
        mat[2, 1] = np.sin(rx)
        mat[2, 2] = np.cos(rx)
        points = np.matmul(points, mat)
    if ry != 0:
        mat = np.zeros((4, 4))
        mat[1, 1] = 1
        mat[3, 3] = 1
        mat[0, 0] = np.cos(ry)
        mat[0, 2] = np.sin(ry)
        mat[2, 0] = -np.sin(ry)
        mat[2, 2] = np.cos(ry)
        points = np.matmul(points, mat)
    if rz != 0:
        mat = np.zeros((4, 4))
        mat[2, 2] = 1
        mat[3, 3] = 1
        mat[0, 0] = np.cos(rz)
        mat[0, 1] = -np.sin(rz)
        mat[1, 0] = np.sin(rz)
        mat[1, 1] = np.cos(rz)
        points = np.matmul(points, mat)
    return points[:, 0:3]

def box_transform(boxes_corner, tx, ty, tz, r=0):
    # boxes_corner (N, 8, 3)
    for idx in range(len(boxes_corner)):
        boxes_corner[idx] = point_transform(boxes_corner[idx], tx, ty, tz, rz=r)
    return boxes_corner

def cal_iou2d(box1_corner, box2_corner):
    box1_corner = np.reshape(box1_corner, [4, 2])
    box2_corner = np.reshape(box2_corner, [4, 2])
    box1_corner = ((config.W, config.H)-(box1_corner - (config.xrange[0], config.yrange[0])) / (config.vw, config.vh)).astype(np.int32)
    box2_corner = ((config.W, config.H)-(box2_corner - (config.xrange[0], config.yrange[0])) / (config.vw, config.vh)).astype(np.int32)

    buf1 = np.zeros((config.H, config.W, 3))
    buf2 = np.zeros((config.H, config.W, 3))
    buf1 = cv2.fillConvexPoly(buf1, box1_corner, color=(1,1,1))[..., 0]
    buf2 = cv2.fillConvexPoly(buf2, box2_corner, color=(1,1,1))[..., 0]

    indiv = np.sum(np.absolute(buf1-buf2))
    share = np.sum((buf1 + buf2) == 2)
    if indiv == 0:
        return 0.0 # when target is out of bound
    return share / (indiv + share)

def aug_data(lidar, gt_box3d_corner):
    np.random.seed()

    choice = np.random.randint(1, 10)

    if choice >= 7:
        for idx in range(len(gt_box3d_corner)):
            # TODO: precisely gather the point
            is_collision = True
            _count = 0
            while is_collision and _count < 100:
                t_rz = np.random.uniform(-np.pi / 10, np.pi / 10)
                t_x = np.random.normal()
                t_y = np.random.normal()
                t_z = np.random.normal()

                # check collision
                tmp = box_transform(
                    gt_box3d_corner[[idx]], t_x, t_y, t_z, t_rz)
                is_collision = False
                for idy in range(idx):
                    iou = cal_iou2d(tmp[0, :4, :2], gt_box3d_corner[idy, :4, :2])
                    if iou > 0:
                        is_collision = True
                        _count += 1
                        break
            if not is_collision:
                box_corner = gt_box3d_corner[idx]
                minx = np.min(box_corner[:, 0])
                miny = np.min(box_corner[:, 1])
                minz = np.min(box_corner[:, 2])
                maxx = np.max(box_corner[:, 0])
                maxy = np.max(box_corner[:, 1])
                maxz = np.max(box_corner[:, 2])
                bound_x = np.logical_and(
                    lidar[:, 0] >= minx, lidar[:, 0] <= maxx)
                bound_y = np.logical_and(
                    lidar[:, 1] >= miny, lidar[:, 1] <= maxy)
                bound_z = np.logical_and(
                    lidar[:, 2] >= minz, lidar[:, 2] <= maxz)
                bound_box = np.logical_and(
                    np.logical_and(bound_x, bound_y), bound_z)
                lidar[bound_box, 0:3] = point_transform(
                    lidar[bound_box, 0:3], t_x, t_y, t_z, rz=t_rz)
                gt_box3d_corner[idx] = box_transform(
                    gt_box3d_corner[[idx]], t_x, t_y, t_z, t_rz)

        gt_box3d = gt_box3d_corner

    elif choice < 7 and choice >= 4:
        # global rotation
        angle = np.random.uniform(-np.pi / 4, np.pi / 4)
        lidar[:, 0:3] = point_transform(lidar[:, 0:3], 0, 0, 0, rz=angle)
        gt_box3d = box_transform(gt_box3d_corner, 0, 0, 0, r=angle)

    else:
        # global scaling
        factor = np.random.uniform(0.95, 1.05)
        lidar[:, 0:3] = lidar[:, 0:3] * factor
        gt_box3d = gt_box3d_corner * factor

    return lidar, gt_box3d

# Utilities

def get_filtered_lidar(lidar, boxes3d=None):

    pxs = lidar[:, 0]
    pys = lidar[:, 1]
    pzs = lidar[:, 2]

    filter_x = np.where((pxs >= config.xrange[0]) & (pxs < config.xrange[1]))[0]
    filter_y = np.where((pys >= config.yrange[0]) & (pys < config.yrange[1]))[0]
    filter_z = np.where((pzs >= config.zrange[0]) & (pzs < config.zrange[1]))[0]
    filter_xy = np.intersect1d(filter_x, filter_y)
    filter_xyz = np.intersect1d(filter_xy, filter_z)

    if boxes3d is not None:
        box_x = (boxes3d[:, :, 0] >= config.xrange[0]) & (boxes3d[:, :, 0] < config.xrange[1])
        box_y = (boxes3d[:, :, 1] >= config.yrange[0]) & (boxes3d[:, :, 1] < config.yrange[1])
        box_z = (boxes3d[:, :, 2] >= config.zrange[0]) & (boxes3d[:, :, 2] < config.zrange[1])
        box_xyz = np.sum(box_x & box_y & box_z,axis=1)

        return lidar[filter_xyz], boxes3d[box_xyz>0]

    return lidar[filter_xyz]

def lidar_to_bev(lidar):

    X0, Xn = 0, config.W
    Y0, Yn = 0, config.H
    Z0, Zn = 0, config.D

    width  = Yn - Y0
    height   = Xn - X0
    channel = Zn - Z0  + 2

    pxs = lidar[:, 0]
    pys = lidar[:, 1]
    pzs = lidar[:, 2]
    prs = lidar[:, 3]

    qxs=((pxs-config.xrange[0])/config.vw).astype(np.int32)
    qys=((pys-config.yrange[0])/config.vh).astype(np.int32)
    qzs=((pzs-config.zrange[0])/config.vd).astype(np.int32)

    print('height,width,channel=%d,%d,%d' % (height, width, channel))
    top = np.zeros(shape=(height, width, channel), dtype=np.float32)
    mask = np.ones(shape=(height, width, channel-1), dtype=np.float32) * -5

    for i in range(len(pxs)):
        top[-qxs[i], -qys[i], -1] = 1 + top[-qxs[i], -qys[i], -1]
        if pzs[i] > mask[-qxs[i], -qys[i], qzs[i]]:
            top[-qxs[i], -qys[i], qzs[i]] = max(0, pzs[i]-config.zrange[0])
            mask[-qxs[i], -qys[i], qzs[i]]=pzs[i]
        if pzs[i] > mask[-qxs[i], -qys[i], -1]:
            mask[-qxs[i], -qys[i], -1] = pzs[i]
            top[-qxs[i], -qys[i], -2] = prs[i]

    top[:, :, -1] = np.log(top[:, :, -1]+1)/math.log(64)

    if 1:
        # top_image = np.sum(top[:,:,:-1],axis=2)
        density_image = top[:, :, -1]
        density_image = density_image-np.min(density_image)
        density_image = (density_image/np.max(density_image)*255).astype(np.uint8)
        # top_image = np.dstack((top_image, top_image, top_image)).astype(np.uint8)

    return top, density_image

def draw_lidar(lidar, is_grid=False, is_axis=True, is_top_region=True, fig=None):

    pxs = lidar[:, 0]
    pys = lidar[:, 1]
    pzs = lidar[:, 2]
    prs = lidar[:, 3]

    if fig is None: fig = mlab.figure(figure=None, bgcolor=(0,0,0), fgcolor=None, engine=None, size=(1000, 500))

    mlab.points3d(
        pxs, pys, pzs, prs,
        mode='point',  # 'point'  'sphere'
        colormap='gnuplot',  # 'bone',  #'spectral',  #'copper',
        scale_factor=1,
        figure=fig)

    # draw grid
    if is_grid:
        mlab.points3d(0, 0, 0, color=(1,1,1), mode='sphere', scale_factor=0.2)

        for y in np.arange(-50,50,1):
            x1,y1,z1 = -50, y, 0
            x2,y2,z2 =  50, y, 0
            mlab.plot3d([x1, x2], [y1, y2], [z1,z2], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)

        for x in np.arange(-50,50,1):
            x1,y1,z1 = x,-50, 0
            x2,y2,z2 = x, 50, 0
            mlab.plot3d([x1, x2], [y1, y2], [z1,z2], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)

    # draw axis
    if is_grid:
        mlab.points3d(0, 0, 0, color=(1,1,1), mode='sphere', scale_factor=0.2)

        axes=np.array([
            [2.,0.,0.,0.],
            [0.,2.,0.,0.],
            [0.,0.,2.,0.],
        ],dtype=np.float64)
        fov=np.array([  ##<todo> : now is 45 deg. use actual setting later ...
            [20., 20., 0.,0.],
            [20.,-20., 0.,0.],
        ],dtype=np.float64)

        mlab.plot3d([0, axes[0,0]], [0, axes[0,1]], [0, axes[0,2]], color=(1,0,0), tube_radius=None, figure=fig)
        mlab.plot3d([0, axes[1,0]], [0, axes[1,1]], [0, axes[1,2]], color=(0,1,0), tube_radius=None, figure=fig)
        mlab.plot3d([0, axes[2,0]], [0, axes[2,1]], [0, axes[2,2]], color=(0,0,1), tube_radius=None, figure=fig)
        mlab.plot3d([0, fov[0,0]], [0, fov[0,1]], [0, fov[0,2]], color=(1,1,1), tube_radius=None, line_width=1, figure=fig)
        mlab.plot3d([0, fov[1,0]], [0, fov[1,1]], [0, fov[1,2]], color=(1,1,1), tube_radius=None, line_width=1, figure=fig)

    # draw top_image feature area
    if is_top_region:
        x1 = config.xrange[0]
        x2 = config.xrange[1]
        y1 = config.yrange[0]
        y2 = config.yrange[1]
        mlab.plot3d([x1, x1], [y1, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)
        mlab.plot3d([x2, x2], [y1, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)
        mlab.plot3d([x1, x2], [y1, y1], [0,0], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)
        mlab.plot3d([x1, x2], [y2, y2], [0,0], color=(0.5,0.5,0.5), tube_radius=None, line_width=1, figure=fig)

    mlab.orientation_axes()
    mlab.view(azimuth=180,elevation=None,distance=50,focalpoint=[ 12.0909996 , -1.04700089, -2.03249991])#2.0909996 , -1.04700089, -2.03249991

    return fig

def draw_gt_boxes3d(gt_boxes3d, fig, color=(1,0,0), line_width=2):

    num = len(gt_boxes3d)
    for n in range(num):
        b = gt_boxes3d[n]

        for k in range(0,4):

            i,j=k,(k+1)%4
            mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig)

            i,j=k+4,(k+3)%4 + 4
            mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig)

            i,j=k,k+4
            mlab.plot3d([b[i,0], b[j,0]], [b[i,1], b[j,1]], [b[i,2], b[j,2]], color=color, tube_radius=None, line_width=line_width, figure=fig)

    mlab.view(azimuth=180,elevation=None,distance=50,focalpoint=[ 12.0909996 , -1.04700089, -2.03249991])#2.0909996 , -1.04700089, -2.03249991

def project_velo2rgb(velo,calib):
    T=np.zeros([4,4],dtype=np.float32)
    T[:3,:]=calib['Tr_velo2cam']
    T[3,3]=1
    R=np.zeros([4,4],dtype=np.float32)
    R[:3,:3]=calib['R0']
    R[3,3]=1
    num=len(velo)
    projections = np.zeros((num,8,2),  dtype=np.int32)
    for i in range(len(velo)):
        box3d=np.ones([8,4],dtype=np.float32)
        box3d[:,:3]=velo[i]
        M=np.dot(calib['P2'],R)
        M=np.dot(M,T)
        box2d=np.dot(M,box3d.T)
        box2d=box2d[:2,:].T/box2d[2,:].reshape(8,1)
        projections[i] = box2d
    return projections

def draw_rgb_projections(image, projections, color=(255,255,255), thickness=2, darker=1):

    img = image.copy()*darker
    num=len(projections)
    forward_color=(255,255,0)
    for n in range(num):
        qs = projections[n]
        for k in range(0,4):
            i,j=k,(k+1)%4

            cv2.line(img, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

            i,j=k+4,(k+1)%4 + 4
            cv2.line(img, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

            i,j=k,k+4
            cv2.line(img, (qs[i,0],qs[i,1]), (qs[j,0],qs[j,1]), color, thickness, cv2.LINE_AA)

        cv2.line(img, (qs[3,0],qs[3,1]), (qs[7,0],qs[7,1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[7,0],qs[7,1]), (qs[6,0],qs[6,1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[6,0],qs[6,1]), (qs[2,0],qs[2,1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[2,0],qs[2,1]), (qs[3,0],qs[3,1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[3,0],qs[3,1]), (qs[6,0],qs[6,1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[2,0],qs[2,1]), (qs[7,0],qs[7,1]), forward_color, thickness, cv2.LINE_AA)

    return img

def _quantize_coords(x, y):
    xx = config.H - int((y - config.yrange[0]) / config.vh)
    yy = config.W - int((x - config.xrange[0]) / config.vw)
    return xx, yy

def  draw_polygons(image, polygons,color=(255,255,255), thickness=1, darken=1):

    img = image.copy() * darken
    for polygon in polygons:
        tup0, tup1, tup2, tup3 = [_quantize_coords(*tup) for tup in polygon]
        cv2.line(img, tup0, tup1, color, thickness, cv2.LINE_AA)
        cv2.line(img, tup1, tup2, color, thickness, cv2.LINE_AA)
        cv2.line(img, tup2, tup3, color, thickness, cv2.LINE_AA)
        cv2.line(img, tup3, tup0, color, thickness, cv2.LINE_AA)
    return img

def draw_rects(image, rects, color=(255,255,255), thickness=1, darken=1):

    img = image.copy() * darken
    for rect in rects:
        tup0,tup1 = [_quantize_coords(*tup) for tup in list(zip(rect[0::2], rect[1::2]))]
        cv2.rectangle(img, tup0, tup1, color, thickness, cv2.LINE_AA)
    return img

def load_kitti_calib(calib_file):
    """
    load projection matrix
    """
    with open(calib_file) as fi:
        lines = fi.readlines()
        assert (len(lines) == 8)

    obj = lines[0].strip().split(' ')[1:]
    P0 = np.array(obj, dtype=np.float32)
    obj = lines[1].strip().split(' ')[1:]
    P1 = np.array(obj, dtype=np.float32)
    obj = lines[2].strip().split(' ')[1:]
    P2 = np.array(obj, dtype=np.float32)
    obj = lines[3].strip().split(' ')[1:]
    P3 = np.array(obj, dtype=np.float32)
    obj = lines[4].strip().split(' ')[1:]
    R0 = np.array(obj, dtype=np.float32)
    obj = lines[5].strip().split(' ')[1:]
    Tr_velo_to_cam = np.array(obj, dtype=np.float32)
    obj = lines[6].strip().split(' ')[1:]
    Tr_imu_to_velo = np.array(obj, dtype=np.float32)

    return {'P2': P2.reshape(3, 4),
            'R0': R0.reshape(3, 3),
            'Tr_velo2cam': Tr_velo_to_cam.reshape(3, 4)}

def angle_in_limit(angle):
    # To limit the angle in -pi/2 - pi/2
    limit_degree = 5
    while angle >= np.pi / 2:
        angle -= np.pi
    while angle < -np.pi / 2:
        angle += np.pi
    if abs(angle + np.pi / 2) < limit_degree / 180 * np.pi:
        angle = np.pi / 2
    return angle

def box3d_cam_to_velo(box3d, Tr):

    def project_cam2velo(cam, Tr):
        T = np.zeros([4, 4], dtype=np.float32)
        T[:3, :] = Tr
        T[3, 3] = 1
        T_inv = np.linalg.inv(T)
        lidar_loc_ = np.dot(T_inv, cam)
        lidar_loc = lidar_loc_[:3]
        return lidar_loc.reshape(1, 3)

    def ry_to_rz(ry):
        angle = -ry - np.pi / 2

        if angle >= np.pi:
            angle -= np.pi
        if angle < -np.pi:
            angle = 2*np.pi + angle

        return angle

    h,w,l,tx,ty,tz,ry = [float(i) for i in box3d]
    cam = np.ones([4, 1])
    cam[0] = tx
    cam[1] = ty
    cam[2] = tz
    t_lidar = project_cam2velo(cam, Tr)

    Box = np.array([[-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2],
                    [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2],
                    [0, 0, 0, 0, h, h, h, h]])

    rz = ry_to_rz(ry)

    rotMat = np.array([
        [np.cos(rz), -np.sin(rz), 0.0],
        [np.sin(rz), np.cos(rz), 0.0],
        [0.0, 0.0, 1.0]])

    velo_box = np.dot(rotMat, Box)

    cornerPosInVelo = velo_box + np.tile(t_lidar, (8, 1)).T

    box3d_corner = cornerPosInVelo.transpose()

    return box3d_corner.astype(np.float32)

def anchors_center_to_corner(anchors):
    N = anchors.shape[0]
    anchor_corner = np.zeros((N, 4, 2))
    for i in range(N):
        anchor = anchors[i]
        translation = anchor[0:3]
        h, w, l = anchor[3:6]
        rz = anchor[-1]
        Box = np.array([
            [-l / 2, -l / 2, l / 2, l / 2], \
            [w / 2, -w / 2, -w / 2, w / 2]])
        # re-create 3D bounding box in velodyne coordinate system
        rotMat = np.array([
            [np.cos(rz), -np.sin(rz)],
            [np.sin(rz), np.cos(rz)]])
        velo_box = np.dot(rotMat, Box)
        cornerPosInVelo = velo_box + np.tile(translation[:2], (4, 1)).T
        box2d = cornerPosInVelo.transpose()
        anchor_corner[i] = box2d
    return anchor_corner

def corner_to_standup_box2d_batch(boxes_corner):
    # (N, 4, 2) -> (N, 4) x1, y1, x2, y2
    N = boxes_corner.shape[0]
    standup_boxes2d = np.zeros((N, 4))
    standup_boxes2d[:, 0] = np.min(boxes_corner[:, :, 0], axis=1)
    standup_boxes2d[:, 1] = np.min(boxes_corner[:, :, 1], axis=1)
    standup_boxes2d[:, 2] = np.max(boxes_corner[:, :, 0], axis=1)
    standup_boxes2d[:, 3] = np.max(boxes_corner[:, :, 1], axis=1)
    return standup_boxes2d

def box3d_corner_to_center_batch(box3d_corner):
    # (N, 8, 3) -> (N, 7)
    assert box3d_corner.ndim == 3
    batch_size = box3d_corner.shape[0]

    xyz = np.mean(box3d_corner[:, :4, :], axis=1)

    h = abs(np.mean(box3d_corner[:, 4:, 2] - box3d_corner[:, :4, 2], axis=1, keepdims=True))

    w = (np.sqrt(np.sum((box3d_corner[:, 0, [0, 1]] - box3d_corner[:, 1, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 2, [0, 1]] - box3d_corner[:, 3, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 4, [0, 1]] - box3d_corner[:, 5, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 6, [0, 1]] - box3d_corner[:, 7, [0, 1]]) ** 2, axis=1, keepdims=True))) / 4

    l = (np.sqrt(np.sum((box3d_corner[:, 0, [0, 1]] - box3d_corner[:, 3, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 1, [0, 1]] - box3d_corner[:, 2, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 4, [0, 1]] - box3d_corner[:, 7, [0, 1]]) ** 2, axis=1, keepdims=True)) +
         np.sqrt(np.sum((box3d_corner[:, 5, [0, 1]] - box3d_corner[:, 6, [0, 1]]) ** 2, axis=1, keepdims=True))) / 4

    theta = (np.arctan2(box3d_corner[:, 2, 1] - box3d_corner[:, 1, 1],
                        box3d_corner[:, 2, 0] - box3d_corner[:, 1, 0]) +
             np.arctan2(box3d_corner[:, 3, 1] - box3d_corner[:, 0, 1],
                        box3d_corner[:, 3, 0] - box3d_corner[:, 0, 0]) +
             np.arctan2(box3d_corner[:, 2, 0] - box3d_corner[:, 3, 0],
                        box3d_corner[:, 3, 1] - box3d_corner[:, 2, 1]) +
             np.arctan2(box3d_corner[:, 1, 0] - box3d_corner[:, 0, 0],
                        box3d_corner[:, 0, 1] - box3d_corner[:, 1, 1]))[:, np.newaxis] / 4

    return np.concatenate([xyz, h, w, l, theta], axis=1).reshape(batch_size, 7)

def get_anchor3d(anchors):
    num = anchors.shape[0]
    anchors3d = np.zeros((num,8,3))
    anchors3d[:, :4, :2] = anchors
    anchors3d[:, :, 2] = config.z_a
    anchors3d[:, 4:, :2] = anchors
    anchors3d[:, 4:, 2] = config.z_a + config.h_a
    return anchors3d

def load_kitti_label(label_file, Tr):

    with open(label_file, 'r') as f:
        lines = f.readlines()

    gt_boxes3d_corner = []

    num_obj = len(lines)

    for j in range(num_obj):
        obj = lines[j].strip().split(' ')

        obj_class = obj[0].strip()
        if obj_class not in config.class_list:
            continue

        box3d_corner = box3d_cam_to_velo(obj[8:], Tr)

        gt_boxes3d_corner.append(box3d_corner)

    gt_boxes3d_corner = np.array(gt_boxes3d_corner).reshape(-1, 8, 3)

    return gt_boxes3d_corner

def test():
    import os
    import glob
    import matplotlib.pyplot as plt

    lidar_path = os.path.join('./data/KITTI/training', "crop/")
    image_path = os.path.join('./data/KITTI/training', "image_2/")
    calib_path = os.path.join('./data/KITTI/training', "calib/")
    label_path = os.path.join('./data/KITTI/training', "label_2/")

    file=[i.strip().split('/')[-1][:-4] for i in sorted(os.listdir(label_path))]

    i = 2600

    lidar_file = lidar_path + '/' + file[i] + '.bin'
    calib_file = calib_path + '/' + file[i] + '.txt'
    label_file = label_path + '/' + file[i] + '.txt'
    image_file = image_path + '/' + file[i] + '.png'

    image = cv2.imread(image_file)
    print("Processing: ", lidar_file)
    lidar = np.fromfile(lidar_file, dtype=np.float32)
    lidar = lidar.reshape((-1, 4))

    calib = load_kitti_calib(calib_file)
    gt_box3d = load_kitti_label(label_file, calib['Tr_velo2cam'])

    # augmentation
    #lidar, gt_box3d = aug_data(lidar, gt_box3d)

    # filtering
    lidar, gt_box3d = get_filtered_lidar(lidar, gt_box3d)

    # view in point cloud

    # fig = draw_lidar(lidar, is_grid=False, is_top_region=True)
    # draw_gt_boxes3d(gt_boxes3d=gt_box3d, fig=fig)
    # mlab.show()

    # view in image

    # gt_3dTo2D = project_velo2rgb(gt_box3d, calib)
    # img_with_box = draw_rgb_projections(image,gt_3dTo2D, color=(0,0,255),thickness=1)
    # plt.imshow(img_with_box[:,:,[2,1,0]])
    # plt.show()

    # view in bird-eye view

    top_new, density_image=lidar_to_bev(lidar)
    # gt_box3d_top = corner_to_standup_box2d_batch(gt_box3d)
    # density_with_box = draw_rects(density_image,gt_box3d_top)
    density_with_box = draw_polygons(density_image,gt_box3d[:,:4,:2])
    plt.imshow(density_with_box,cmap='gray')
    plt.show()

# Train Network

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        init.xavier_uniform(m.weight.data)
        m.bias.data.zero_()

def detection_collate(batch):
    voxel_features = []
    voxel_coords = []
    pos_equal_one = []
    neg_equal_one = []
    targets = []

    images = []
    calibs = []
    ids = []
    for i, sample in enumerate(batch):
        voxel_features.append(sample[0])

        voxel_coords.append(
            np.pad(sample[1], ((0, 0), (1, 0)),
                mode='constant', constant_values=i))

        pos_equal_one.append(sample[2])
        neg_equal_one.append(sample[3])
        targets.append(sample[4])

        images.append(sample[5])
        calibs.append(sample[6])
        ids.append(sample[7])
    return np.concatenate(voxel_features), \
           np.concatenate(voxel_coords), \
           np.array(pos_equal_one),\
           np.array(neg_equal_one),\
           np.array(targets),\
           images, calibs, ids

# Test utilities

def delta_to_boxes3d(deltas, anchors):
    # Input:
    #   deltas: (N, w, l, 14)
    #   feature_map_shape: (w, l)
    #   anchors: (w, l, 2, 7)

    # Ouput:
    #   boxes3d: (N, w*l*2, 7)
    N = deltas.shape[0]
    deltas = deltas.view(N, -1, 7)
    anchors = torch.FloatTensor(anchors)
    boxes3d = torch.zeros_like(deltas)

    if deltas.is_cuda:
        anchors = anchors.cuda()
        boxes3d = boxes3d.cuda()

    anchors_reshaped = anchors.view(-1, 7)

    anchors_d = torch.sqrt(anchors_reshaped[:, 4]**2 + anchors_reshaped[:, 5]**2)

    anchors_d = anchors_d.repeat(N, 2, 1).transpose(1,2)
    anchors_reshaped = anchors_reshaped.repeat(N, 1, 1)

    boxes3d[..., [0, 1]] = torch.mul(deltas[..., [0, 1]], anchors_d) + anchors_reshaped[..., [0, 1]]
    boxes3d[..., [2]] = torch.mul(deltas[..., [2]], anchors_reshaped[...,[3]]) + anchors_reshaped[..., [2]]

    boxes3d[..., [3, 4, 5]] = torch.exp(
        deltas[..., [3, 4, 5]]) * anchors_reshaped[..., [3, 4, 5]]

    boxes3d[..., 6] = deltas[..., 6] + anchors_reshaped[..., 6]

    return boxes3d

def detection_collate(batch):
    lidars = []
    images = []
    calibs = []

    targets = []
    pos_equal_ones=[]
    ids = []
    for i, sample in enumerate(batch):
        lidars.append(sample[0])
        images.append(sample[1])
        calibs.append(sample[2])
        targets.append(sample[3])
        pos_equal_ones.append(sample[4])
        ids.append(sample[5])
    return lidars,images,calibs,\
           torch.cuda.FloatTensor(np.array(targets)), \
           torch.cuda.FloatTensor(np.array(pos_equal_ones)),\
           ids

def box3d_center_to_corner_batch(boxes_center):
    # (N, 7) -> (N, 8, 3)
    N = boxes_center.shape[0]
    ret = torch.zeros((N, 8, 3))
    if boxes_center.is_cuda:
        ret = ret.cuda()

    for i in range(N):
        box = boxes_center[i]
        translation = box[0:3]
        size = box[3:6]
        rotation = [0, 0, box[-1]]

        h, w, l = size[0], size[1], size[2]
        trackletBox = torch.FloatTensor([  # in velodyne coordinates around zero point and without orientation yet
            [-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2], \
            [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2], \
            [0, 0, 0, 0, h, h, h, h]])
        if boxes_center.is_cuda:
            trackletBox = trackletBox.cuda()
        # re-create 3D bounding box in velodyne coordinate system
        yaw = rotation[2]
        rotMat = torch.FloatTensor([
            [np.cos(yaw), -np.sin(yaw), 0.0],
            [np.sin(yaw), np.cos(yaw), 0.0],
            [0.0, 0.0, 1.0]])
        if boxes_center.is_cuda:
            rotMat = rotMat.cuda()

        cornerPosInVelo = torch.mm(rotMat, trackletBox) + translation.repeat(8, 1).t()
        box3d = cornerPosInVelo.transpose(0,1)
        ret[i] = box3d

    return ret

def box3d_corner_to_top_batch(boxes3d, use_min_rect=True):
    # [N,8,3] -> [N,4,2] -> [N,8]
    box3d_top=[]

    num = len(boxes3d)
    for n in range(num):
        b  = boxes3d[n]
        x0 = b[0,0]
        y0 = b[0,1]
        x1 = b[1,0]
        y1 = b[1,1]
        x2 = b[2,0]
        y2 = b[2,1]
        x3 = b[3,0]
        y3 = b[3,1]
        box3d_top.append([x0,y0,x1,y1,x2,y2,x3,y3])

    if use_min_rect:
        box8pts = torch.FloatTensor(np.array(box3d_top))
        if boxes3d.is_cuda:
            box8pts = box8pts.cuda()
        min_rects = torch.zeros((box8pts.shape[0], 4))
        if boxes3d.is_cuda:
            min_rects = min_rects.cuda()
        # calculate minimum rectangle
        min_rects[:, 0] = torch.min(box8pts[:, [0, 2, 4, 6]], dim=1)[0]
        min_rects[:, 1] = torch.min(box8pts[:, [1, 3, 5, 7]], dim=1)[0]
        min_rects[:, 2] = torch.max(box8pts[:, [0, 2, 4, 6]], dim=1)[0]
        min_rects[:, 3] = torch.max(box8pts[:, [1, 3, 5, 7]], dim=1)[0]
        return min_rects

    return box3d_top

def draw_boxes(reg, prob, images, calibs, ids, tag):
    prob = prob.view(config.N, -1)
    batch_boxes3d = delta_to_boxes3d(reg, config.anchors)
    mask = torch.gt(prob, config.score_threshold)
    mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)

    for batch_id in range(config.N):
        boxes3d = torch.masked_select(batch_boxes3d[batch_id], mask_reg[batch_id]).view(-1, 7)
        scores = torch.masked_select(prob[batch_id], mask[batch_id])

        image = images[batch_id]
        calib = calibs[batch_id]
        id = ids[batch_id]

        if len(boxes3d) != 0:

            boxes3d_corner = box3d_center_to_corner_batch(boxes3d)
            boxes2d = box3d_corner_to_top_batch(boxes3d_corner)
            boxes2d_score = torch.cat((boxes2d, scores.unsqueeze(1)), dim=1)

            # NMS
            keep = pth_nms(boxes2d_score, config.nms_threshold)
            boxes3d_corner_keep = boxes3d_corner[keep]
            print("No. %d objects detected" % len(boxes3d_corner_keep))

            rgb_2D = project_velo2rgb(boxes3d_corner_keep, calib)
            img_with_box = draw_rgb_projections(image, rgb_2D, color=(0, 0, 255), thickness=1)
            cv2.imwrite('results/%s_%s.png' % (id,tag), img_with_box)

        else:
            cv2.imwrite('results/%s_%s.png' % (id,tag), image)
            print("No objects detected")

def voxelnet_unit():

    torch.backends.cudnn.enabled = True

    # dataset
    dataset = A2D2Dataset(cfg=config, root='./data/KITTI', set='train')
    data_loader = data.DataLoader(dataset, batch_size=config.N, num_workers=4, collate_fn=detection_collate, shuffle=True, \
                                  pin_memory=False)

    # network
    net = VoxelNet()
    net.cuda()

    net.train()

    # initialization
    print('Initializing weights...')
    net.apply(weights_init)

    # define optimizer
    optimizer = optim.SGD(net.parameters(), lr=0.01)

    # define loss function
    criterion = VoxelLoss(alpha=1.5, beta=1)

    # training process
    batch_iterator = None
    epoch_size = len(dataset) // config.N
    print('Epoch size', epoch_size)
    for iteration in range(10000):
        if (not batch_iterator) or (iteration % epoch_size == 0):
            # create batch iterator
            batch_iterator = iter(data_loader)

        voxel_features, voxel_coords, pos_equal_one, neg_equal_one, targets, images, calibs, ids = next(batch_iterator)

        # wrapper to variable
        voxel_features = Variable(torch.cuda.FloatTensor(voxel_features))
        pos_equal_one = Variable(torch.cuda.FloatTensor(pos_equal_one))
        neg_equal_one = Variable(torch.cuda.FloatTensor(neg_equal_one))
        targets = Variable(torch.cuda.FloatTensor(targets))

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward
        t0 = time.time()
        psm, rm = net(voxel_features, voxel_coords)

        # calculate loss
        conf_loss, reg_loss = criterion(rm, psm, pos_equal_one, neg_equal_one, targets)
        loss = conf_loss + reg_loss

        # backward
        loss.backward()
        optimizer.step()

        t1 = time.time()

        print('Timer: %.4f sec.' % (t1 - t0))
        print('iter ' + repr(iteration) + ' || Loss: %.4f || Conf Loss: %.4f || Loc Loss: %.4f' % \
              (loss.data[0], conf_loss.data[0], reg_loss.data[0]))

        # visualization
        # draw_boxes(rm, psm, ids, images, calibs, 'pred')
        draw_boxes(targets.data, pos_equal_one.data, images, calibs, ids, 'true')