#-*- encoding:utf-8 -*-
'''
Created on Dec 1, 2014
@author: letian
'''
import networkx as nx
from Segmentation import Segmentation
import numpy as np
import math

class TextRank4Sentence(object):

    def __init__(self, stop_words_file = None, delimiters='?!;？！。；…\n'):
        '''
        `stop_words_file`：默认值为None，此时内部停止词表为空；可以设置为文件路径（字符串），将从停止词文件中提取停止词。
        `delimiters`：默认值是`'?!;？！。；…\n'`，用来将文本拆分为句子。

        self.sentences：由句子组成的列表。
        self.words_no_filter：对sentences中每个句子分词而得到的两级列表。
        self.words_no_stop_words：去掉words_no_filter中的停止词而得到的两级列表。
        self.words_all_filters：保留words_no_stop_words中指定词性的单词而得到的两级列表。
        '''
        self.seg = Segmentation(stop_words_file=stop_words_file, delimiters=delimiters)

        self.sentences = None
        self.words_no_filter = None     # 2维列表
        self.words_no_stop_words = None
        self.words_all_filters = None

        self.graph = None
        self.key_sentences = None

    def train(self, text, lower = False, speech_tag_filter=True,
              source = 'no_stop_words', sim_func = 'standard'):
        '''
       `text`：文本内容，字符串。
        `lower`：是否将文本转换为小写。默认为False。
        `speech_tag_filter`：若值为True，将调用内部的词性列表来过滤生成words_all_filters。
                        若值为False，words_all_filters与words_no_stop_words相同。
        `source`：选择使用words_no_filter, words_no_stop_words, words_all_filters中的哪一个来生成句子之间的相似度。
                默认值为`'all_filters'`，可选值为`'no_filter', 'no_stop_words', 'all_filters'`。
        `sim_func`： 指定计算句子相似度的函数。当前只有一个函数，对应默认值`standard`。
        '''

        self.key_sentences = []

        (self.sentences, self.words_no_filter, self.words_no_stop_words, self.words_all_filters) = self.seg.segment(text=text,
                                                                                lower=lower,
                                                                                speech_tag_filter=speech_tag_filter);
        # -

        # print self.sentences

        if source == 'no_filter':
            source = self.words_no_filter
        elif source == 'all_filters':
            source = self.words_all_filters
        else:
            source = self.words_no_stop_words

        sim_func = self._get_similarity_standard

        sentences_num = len(source)

        self.graph = np.zeros((sentences_num, sentences_num))

        for x in xrange(sentences_num):
            for y in xrange(x, sentences_num):
                similarity = sim_func(source[x], source[y])
                self.graph[x, y] = similarity
                self.graph[y, x] = similarity

#         for x in xrange(sentences_num):
#             row_sum = np.sum(self.graph[x, :])
#             if row_sum > 0:
#                 self.graph[x, :] = self.graph[x, :] / row_sum

        # print self.graph

        nx_graph = nx.from_numpy_matrix(self.graph)
        scores = nx.pagerank(nx_graph) # this is a dict
        sorted_scores = sorted(scores.items(), key = lambda item: item[1], reverse=True)

        # print sorted_scores

        for index, _ in sorted_scores:
            self.key_sentences.append(self.sentences[index])

        # print '\n'.join(self.key_sentences)

    def _get_similarity_standard(self, word_list1, word_list2):
        '''
        默认的用于计算两个句子相似度的函数。
        word_list1, word_list2: 分别代表两个句子，都是由单词组成的列表
        '''
        vector1, vector2 =self._gen_vectors(word_list1, word_list2)

        # print vector1, vector2

        vector3 = [vector1[x]*vector2[x]  for x in xrange(len(vector1))]
        vector4 = [1 for num in vector3 if num > 0.]
        co_occur_num = sum(vector4)

        # print co_occur_num

        if co_occur_num == 0.:
            return 0.

        denominator = math.log(float(len(word_list1))) + math.log(float(len(word_list2))) # 分母

        if denominator == 0.:
            return 0.

        return co_occur_num / denominator

    def _gen_vectors(self, word_list1, word_list2):
        '''
        两个句子转换成两个同样大小向量。可以通过这两个向量来计算两个句子的相似度。
        word_list1, word_list2: 分别代表两个句子，都是由单词组成的列表
        '''
        words = list(set(word_list1 + word_list2))
        vector1 = [float(word_list1.count(word)) for word in words]
        vector2 = [float(word_list2.count(word)) for word in words]
        return vector1, vector2

    def get_key_sentences(self, num = 6, sentence_min_len = 6):
        '''
        获取最重要的num个长度大于等于sentence_min_len的句子用来生成摘要。
        返回列表。
        '''
        result = []
        count = 0
        for sentence in self.key_sentences:
            if count >= num:
                break
            if len(sentence) >= sentence_min_len:
                result.append(sentence)
                count += 1
        return result

if __name__ == '__main__':

    import codecs
    # text = codecs.open('../text/03.txt', 'r', 'utf-8').read()
    text = "这间酒店位于北京东三环，里面摆放很多雕塑，文艺气息十足。答谢宴于晚上8点开始。"
    tr4s = TextRank4Sentence(stop_words_file='../stopword.data')
    tr4s.train(text=text, speech_tag_filter=True, lower=True, source = 'all_filters')
    print '\n'.join(tr4s.get_key_sentences(num=1))

    print '\n'.join(tr4s.sentences)
    for wl in tr4s.words_no_filter:
        print '[', ', \''.join(wl), ']'
    print
    for wl in tr4s.words_no_stop_words:
        print '[', ', \''.join(wl), ']'
    print
    for wl in tr4s.words_all_filters:
        print '[', ', \''.join(wl), ']'


# -*- coding: utf-8 -*-

# IMPORTS
#import matplotlib as mpl
#mpl.use('Agg')
import matplotlib.pyplot as plt

import os
import numpy as np
import json
from skimage import io
from scipy import ndimage
from skimage.io import imread, imsave

from rnaloc import toolbox

# Turn off warnings occuring when saving images
import warnings
warnings.filterwarnings("ignore", message=".*(is a low contrast image)")

def folder_scan_process(folder_root,region_labels,log_msg_callback = None,log_prog_callback=None):
    ''' Scan folders an process when annotation file is present'''

    # Recursive search of specified directory
    folders_proc = []
    for root, dirnames, filenames in os.walk(folder_root):
        for filename in filenames:
            if filename == "annotation.json" :
                folders_proc.append(root)

    for folder_process in folders_proc:
        toolbox.log_message(f'\n\n Processing folder: {folder_process}',callback_fun=log_msg_callback)
        process_folder(folder_process,region_labels,log_msg_callback=log_msg_callback,log_prog_callback=log_prog_callback)

def process_folder(folder_process,region_labels,log_msg_callback = None,log_prog_callback=None):
    '''
    Process folder containing FQ analysis results and an ImJoy annotation.
    '''

    # Open annotations and create masks
    annotation_file = os.path.join(folder_process,"annotation.json")
    data_json, img_size = toolbox.read_annotation(annotation_file)

    # Note that coordinates are exchanged and y flipped
    for feat_idx, feat in enumerate(data_json['features']):
        label = feat['properties']['label']

        if label == region_labels[0]:
            toolbox.log_message(f'Annotation for first region ({region_labels[0]}) found',callback_fun=log_msg_callback)
            cv_pos = np.squeeze(np.asarray(feat['geometry']['coordinates']))
            cv_pos[:,[0, 1]] = cv_pos[:,[1, 0]]
            cv_pos[:,0] = -1*cv_pos[:,0]+img_size[0]
            reg1_mask = toolbox.make_mask(cv_pos,img_size)

        elif label == region_labels[1]:
            toolbox.log_message(f'Annotation for second region ({region_labels[1]}) found',callback_fun=log_msg_callback)
            pl_pos = np.squeeze(np.asarray(feat['geometry']['coordinates']))
            pl_pos[:,[0, 1]] = pl_pos[:,[1, 0]]
            pl_pos[:,0] = -1*pl_pos[:,0]+img_size[0]
            reg2_mask = toolbox.make_mask(pl_pos,img_size)

    ### Make masks and measure distance of P.L. to C.V.

    # Assemble distance map: outside positive, inside negative
    reg1_mask_distTrans_inside = ndimage.distance_transform_edt(reg1_mask)
    reg1_mask_distTrans_outside = ndimage.distance_transform_edt(~reg1_mask.astype(bool))

    reg1_mask_distTrans = np.copy(reg1_mask_distTrans_outside)
    reg1_mask_distTrans[reg1_mask_distTrans_inside>0] = -reg1_mask_distTrans_inside[reg1_mask_distTrans_inside>0]

    # Center of mass of region 2 (portal lobe)
    reg2_com      = np.asarray(ndimage.measurements.center_of_mass(reg2_mask.astype(bool))).astype('int')
    reg2_distReg1 = reg1_mask_distTrans[reg2_com[0],reg2_com[1]]

    # Loop over all FQ result files
    for file in os.listdir(folder_process):
        if '_spots_' in file:

            file_open = os.path.join(folder_process,file)
            toolbox.log_message(f'\n Opening FQ file: {file_open}',callback_fun=log_msg_callback)

            # Get information (path, file name) to save results
            drive, path_and_file = os.path.splitdrive(file_open)
            path, file = os.path.split(path_and_file)
            file_base, ext = os.path.splitext(file)

            path_save = os.path.join(drive,path, 'analysis__exprGradient')
            toolbox.log_message(f'Results will be saved in folder: {path_save}',callback_fun=log_msg_callback)
            if not os.path.isdir(path_save):
                os.makedirs(path_save)

            fq_dict = toolbox.read_FQ_matlab(file_open)
            spots_all = toolbox.get_rna(fq_dict)
            spots_pos = spots_all[:,[16, 17]].astype('int')

            # Open FISH image
            file_FISH_img = os.path.join(folder_process,fq_dict['file_names']['smFISH'])
            toolbox.log_message(f'Reading FISH image: {file_FISH_img}',callback_fun=log_msg_callback)
            img_FISH  = imread(file_FISH_img)

            ## Generate density plots
            name_save = os.path.join(path_save, '_summary_density__' + file_base +  '.png')
            img_density,img_density_outline,img_outline = toolbox.calc_expression_density_plot(fq_dict ,img_size,name_save=name_save,log_msg_callback=log_msg_callback,log_prog_callback=log_prog_callback)

            #Save density image
            imsave(os.path.join(path_save, 'img_density__' + file_base +  '.tif'),img_density)
            imsave(os.path.join(path_save, 'img_density_outline__' + file_base +  '.tif'),img_density_outline)
            imsave(os.path.join(path_save, 'img_outline__' + file_base +  '.tif'),img_outline)

            ### Distance measurements

            # Distance of all RNAs to region 1, RNAs inside the region have negative values
            RNAdist = reg1_mask_distTrans[spots_pos[:,0],spots_pos[:,1]]

            # Renormalize distance map and RNA distances
            reg1_mask_distTrans_norm = np.divide(reg1_mask_distTrans,reg2_distReg1)
            RNAdist_norm = np.divide(RNAdist,reg2_distReg1)

            # Bins for histogram
            RNAdist_norm_max = np.amax(RNAdist_norm)
            RNAdist_norm_min = np.amin(RNAdist_norm)
            bins=np.arange(np.around(RNAdist_norm_min,1)-0.1,RNAdist_norm_max+0.1,0.1)
            width = 0.8 * (bins[1] - bins[0])
            center = (bins[:-1] + bins[1:]) / 2

            # Histogram of RNA distances
            count_RNA, bins = np.histogram(RNAdist_norm, bins=bins,density=False)

            # Renormalize considering how many pixels are really in the actual image
            count_pix = np.diff(list(map(lambda threshold: np.sum(reg1_mask_distTrans_norm <= threshold),bins)))

            # Renormalize RNA counts with respect to area
            count_RNA_normArea = count_RNA/count_pix

            # Renormalize to have sum 1
            hist_sum = np.nansum(count_RNA_normArea)*((bins[1] - bins[0]))
            count_RNA_normSum = np.divide(count_RNA_normArea,hist_sum)

            # Summarize all histograms
            hist_all = np.stack((center,count_RNA_normSum,count_RNA,count_pix),axis=1)

            # Save file with histogram
            np.savetxt(os.path.join(path_save, 'hist_expression__' + file_base +  '.txt'), hist_all, fmt='%f \t %f \t %f \t %f',header='Dist_norm [um]\tCOUNTS_RAW\tCOUNTS_NORM_sum\tPIXEL_COUNTS')

            # Plot results and save figure

            # PLOT ROI and center of mass
            fig1, ax = plt.subplots(3,2,num='dist_enrich')
            fig1.set_size_inches((15,12))

            # Plot image with region of interest and reference point
            img1 = ax[0][0].imshow(img_FISH,cmap="hot")
            plt.sca(ax[0][0])   # set current axis
            plt.title('Region 1 (green) and 2 (blue)')
            ax[0][0].plot(pl_pos[:,1], pl_pos[:,0], '-b')
            ax[0][0].plot(cv_pos[:,1], cv_pos[:,0], '-g')
            ax[0][0].scatter(reg2_com[1],reg2_com[0],color='b')
            ax[0][0].get_xaxis().set_visible(False)
            ax[0][0].get_yaxis().set_visible(False)
            toolbox.colorbar(img1)

            # Plot image with region of interest and reference point
            img1 = ax[0][1].imshow(img_density,cmap="hot")
            plt.sca(ax[0][1])   # set current axis
            plt.title('Region 1 (green) and 2 (blue)')
            ax[0][1].plot(pl_pos[:,1], pl_pos[:,0], '-b')
            ax[0][1].plot(cv_pos[:,1], cv_pos[:,0], '-g')
            ax[0][1].scatter(reg2_com[1],reg2_com[0],color='b')
            ax[0][1].get_xaxis().set_visible(False)
            ax[0][1].get_yaxis().set_visible(False)
            toolbox.colorbar(img1)

            # Plot distance map and pixel distance histogram
            img3 = ax[1][0].imshow(reg1_mask_distTrans_norm,cmap="hot")
            plt.sca(ax[1][0])
            plt.title('Renormalized distance from region 1')
            ax[1][0].plot(pl_pos[:,1], pl_pos[:,0], '-b')
            ax[1][0].plot(cv_pos[:,1], cv_pos[:,0], '-g')
            ax[1][0].scatter(reg2_com[1],reg2_com[0],color='b')
            ax[1][0].get_xaxis().set_visible(False)
            ax[1][0].get_yaxis().set_visible(False)
            toolbox.colorbar(img3)

            ax[1][1].bar(center, count_pix, align='center', width=width)
            ax[1][1].set_xlabel('Distance [pixel]')
            ax[1][1].set_ylabel('# pixel')
            ax[1][1].title.set_text('Histogram of all pixel distances')

            # Plot histograms
            ax[2][0].bar(center, count_RNA, align='center', width=width)
            ax[2][0].set_xlabel('Distance [pixel]')
            ax[2][0].set_ylabel('# RNAs')
            ax[1][0].title.set_text('Histogram without normalization')

            ax[2][1].bar(center, count_RNA_normSum, align='center', width=width)
            ax[2][1].set_xlabel('Normalized distance')
            ax[2][1].set_ylabel('Expression level [a.u.]')
            ax[2][1].title.set_text('Histogram: normalized with area and sum 1')

            fig1.tight_layout(h_pad=0.2)
            plt.draw()

            plt.savefig(os.path.join(path_save, '_summary_gradient_' + file_base +  '.png'),dpi=200)
            plt.close()

    toolbox.log_message(f'Finished processing data!',callback_fun=log_msg_callback)

def process_file(file_open, img_size):
    """
    Analyzes the specified file.
    """

    # Get information (path, file name) to save results
    drive, path_and_file = os.path.splitdrive(file_open)
    path, file = os.path.split(path_and_file)
    file_base, ext = os.path.splitext(file)

    path_save = os.path.join(path, 'analysis_exprdensity')
    if not os.path.isdir(path_save):
        os.makedirs(path_save)

    # Some infos
    print('\n=== Processing file')
    print(file)

    # Open FQ results file
    fq_dict = read_FQ_matlab(file_open)
    spots_all = get_rna(fq_dict)

    # Generate density plots
    print('\n===  Generating density plots')
    img_density,img_density_outline,img_outline = calc_expression_density_plot(fq_dict ,img_size,flag_plot=True)
    plt.figure('density_plt')
    plt.savefig(os.path.join(path_save, 'summary_density__' + file_base +  '.png'),dpi=600)
    plt.close('density_plt')

    # Save density image
    print('-- Saving density plots')

    io.imsave(os.path.join(path_save, 'img_density__' + file_base +  '.tif'),img_density)
    io.imsave(os.path.join(path_save, 'img_density_outline__' + file_base +  '.tif'),img_density_outline)
    io.imsave(os.path.join(path_save, 'img_outline__' + file_base +  '.tif'),img_outline)

    # Select roi
    print('\n=== Select region for calculation of reference point')
    # FROM https://github.com/jdoepfert/roipoly.py
    fig, (ax1) = plt.subplots()
    ax1.imshow(img_density,cmap="hot")
    ax1.get_xaxis().set_visible(False)
    ax1.get_yaxis().set_visible(False)
    MyROI = roipoly(roicolor='g')
    plt.draw()
    plt.close()

    # Analyze distance distribution
    print('==  Analyzing distance distribution')
    ROImask = MyROI.getMask(img_density)
    ROIcom  = ndimage.measurements.center_of_mass(ROImask)
    fq_dict['ref_pos'] = {'com' : ROIcom, 'x':MyROI.allxpoints,'y':MyROI.allypoints}

    hist_all = calc_dist_enrichment(ROIcom,spots_all[:,[16, 17]],img_size,img_density=img_density,flag_plot=True)
    plt.figure('dist_enrich')
    plt.savefig(os.path.join(path_save, 'summary_expgradient__' + file_base +  '.png'),dpi=600)
    plt.close('dist_enrich')

    # Save file with histogram
    np.savetxt(os.path.join(path_save, 'hist_expression__' + file_base +  '.txt'), hist_all, fmt='%10d \t %10d \t %10f \t %10f',header='Dist [um]\tCOUNTS_RAW\tCOUNTS_norm_area\tCOUNTS_NORM_pixel')

    # Save everything to json to be reloaed later if needed
    file_json = os.path.join(path_save, 'fqdict_' + file_base +  '.json')
    with open(file_json,'w') as fp:
        json.dump(fq_dict, fp, cls=NumpyEncoder)


from PyQt5 import QtWidgets, QtGui, QtCore, QtOpenGL, Qt
from OpenGL.GL import *
from OpenGL import GL
from OpenGL.GLU import *

import time
import logging
import numpy as np
from copper import hou
import math
from PIL import Image

import threading
import moderngl

from copper.core.op.op_node import OP_Node
from copper.obj import ObjNode
from copper.ui.utils import clearLayout
from copper.ui.signals import signals
from copper.ui.widgets import PathBarWidget, CollapsableWidget
from copper.ui.panels.base_panel import PathBasedPaneTab
from copper.renderers import Workbench, WorkbenchIPR

from copper.core.vmath import Matrix4, Vector3
from .camera import Camera

from .scene_manager import OGL_Scene_Manager, scene_manager
from .scene_manager.drawable import QuadFS

from pyrr import Matrix44

from .qmodernglwidget import QModernGLWidget

from moderngl_window.opengl.vao import VAO

logger = logging.getLogger(__name__)

class HUD_Info():
    def __init__(self):
        self.fps = 0

class Pixmap_Overlay(QtWidgets.QWidget):
    def __init__(self):
        QtWidgets.QWidget.__init__(self)
        self.setPixmap()
        self.setAttribute(Qt.Qt.WA_NoSystemBackground, True)
        self.setAttribute(Qt.Qt.WA_TransparentForMouseEvents, True)
        self.setAttribute(Qt.Qt.WA_TranslucentBackground, True)

    def setPixmap(self, pixmap=None):
        if pixmap:
            self._pixmap = pixmap
            self.paintEvent = self._paintEvent
        else:
            self.paintEvent = self._paintEventNull

    def _paintEventNull(self, event):
        pass

    def _paintEvent(self, event):
        painter = QtGui.QPainter(self)
        painter.drawPixmap(0, 0, self._pixmap)
        painter.end()

class Signals(QtCore.QObject):
    request_aa_pass = QtCore.pyqtSignal(int)

    def __init__(self, parent=None):
        QtCore.QObject.__init__(self, parent)

from enum import IntEnum

class GeometryViewport(QModernGLWidget):

    class viewType(IntEnum):
        PERSP = 0
        TOP = 1
        BOTTOM = 2
        LEFT = 3
        RIGHT = 4
        FRONT = 5
        BACK = 6

    def __init__(self, parent, panel=None, view_type=0, share_widget=None, scene_manager=None):
        super(GeometryViewport, self).__init__(parent)

        # panel section
        self.panel = panel
        self.setMinimumSize(160, 160)
        self.setSizePolicy(QtWidgets.QSizePolicy( QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding))

        # mouse interactions
        self.orbit_mode = False
        self.old_mouse_x = self.old_mouse_y = 0

        # some viewport stuff
        self.cameras = {
            'persp': Camera(),
            'top': Camera(position=[10,0,0], is_perspective=False),
            'left': Camera(position=[0,0,10], is_perspective=False),
        }

        if view_type == GeometryViewport.viewType.PERSP:
            self.current_camera = 'persp'
        elif view_type == GeometryViewport.viewType.TOP:
            self.current_camera = 'top'
        else:
            self.current_camera = 'left'

        self._show_points = None
        self._show_normals = None
        self._show_hud = None
        self._init_done = False

        self.scene_manager = OGL_Scene_Manager()

        self._width, self._height = None, None

        # offscreen render target / offscreen HUD pixmap
        self.offscreen = None
        self.offscreen2 = None
        self.hud_pixmap = None

        # helpers
        self.m_identity = Matrix44.identity() # just a helper

        # HUD seciton
        self.hud_font = QtGui.QFont("verdana", 8)
        self.hud_info = HUD_Info()
        self.hud_overlay = Pixmap_Overlay()
        self.hud_overlay.setParent(self)
        self.hud_overlay.hide()

        # connect panel signals
        #self.panel.signals.copperNodeModified[OP_Node].connect(self.updateNodeDisplay)

        # connect panel buttons signals
        self.panel.display_options.toggle_points_btn.pressed.connect(self.toggleShowPoints)
        self.panel.display_options.toggle_normals_btn.pressed.connect(self.toggleShowNormals)
        self.panel.display_options.toggle_hud_btn.pressed.connect(self.toggleShowHUD)

        # aa signalling
        self.signals = Signals(self)

        # scene manager signals
        self.scene_manager.signals.geometryUpdated.connect(self.updateGeometryDisplay)

        logger.debug("GeometryViewport widget created")

    def minimumSizeHint(self):
        return QtCore.QSize(200, 200)

    def sizeHint(self):
        return QtCore.QSize(400, 400)

    def drawSceneObjects(self, m_view, m_proj):
        for drawable in self.scene_manager.objects():
            # draw polygons
            drawable.uniformWrite("model", self.m_identity.astype('f4').tobytes())
            drawable.uniformWrite("view", m_view.astype('f4').tobytes())
            drawable.uniformWrite("projection", m_proj.astype('f4').tobytes())
            drawable.draw(show_points = self._show_points)

    @QtCore.pyqtSlot()
    def toggleShowPoints(self):
        self._show_points = not self._show_points
        self.update()

    @QtCore.pyqtSlot()
    def toggleShowNormals(self):
        self._show_normals = not self._show_normals
        self.update()

    @QtCore.pyqtSlot()
    def toggleShowHUD(self):
        self._show_hud = not self._show_hud
        if self._show_hud:
            self.hud_overlay.show()
        else:
            self.hud_overlay.hide()

    @QtCore.pyqtSlot()
    def updateGeometryDisplay(self):
        # Now using quick and dirty hack to check that only geometry nodes changes reflected in scene view
        self.update()

    @QtCore.pyqtSlot()
    def handleRenderedSample(self):
        if self.renderer.image_data.shape[0:2] == (self._width, self._height):
            self.makeCurrent()
            self.offscreen2_diffuse.write(self.renderer.image_data)
            self.update()

    @QtCore.pyqtSlot()
    def handleFinishedSamples(self):
        pass

    def init(self):
        if not self._init_done:
            self.makeCurrent()
            self.ctx.point_size = 2.0
            self.scope = None

            self.buildOffscreen(self.ctx.viewport[2], self.ctx.viewport[3])
            self.scene_manager.init()

            # A fullscreen quad just for rendering one pass to offscreen textures
            self.quad_fs = QuadFS(self)
            self.quad_fs.program['m_model'].write(Matrix44.identity().astype('f4').tobytes())
            self.quad_fs.program['m_view'].write(Matrix44.identity().astype('f4').tobytes())
            self.quad_fs.program['m_proj'].write(Matrix44.orthogonal_projection(-1, 1, 1, -1, 1, 10).astype('f4').tobytes())

            # init renderer
            self.renderer = WorkbenchIPR(scene=None, camera = self.activeCamera)
            self.renderer.setImageUpdateHandler(self.handleRenderedSample)

            self._init_done = True

    def renderHUD(self, aa_pass_num):
        if aa_pass_num == 0:
            # render aa pass independent HUD info
            self.hud_pixmap.fill(QtGui.QColor(0, 0, 0, 0))
            painter = QtGui.QPainter(self.hud_pixmap)
            painter.setPen(QtGui.QColor(255, 128, 16))
            painter.setFont(self.hud_font)
            painter.drawText(10, 20, "FPS: %.1f" % self.hud_info.fps)
            painter.end()
            self.hud_overlay.update()
        else:
            # render aa pass dependent HUD info
            pass

    def resetProgressiveRender(self):
        self._progressive_render_started = False
        self.renderer.resetProgressiveRender()

    def render(self):
        #print("GeometryViewport render")

        self.makeCurrent()
        start_time = time.time()

        m_view = self.activeCamera.getTransform()
        m_proj = self.activeCamera.getProjection()

        # Render the scene to offscreen buffer
        self.offscreen.use()
        self.offscreen.clear(0.0)

        self.ctx.multisample = True
        self.ctx.disable(moderngl.DEPTH_TEST)

        # Render guides

        self.ctx.enable(moderngl.DEPTH_TEST)
        # TODO: we might also want to pass model matrix so we can get different grid orientations instead of rebuilding grid
        self.scene_manager.grid.view.write(m_view.astype('f4').tobytes())
        self.scene_manager.grid.projection.write(m_proj.astype('f4').tobytes())
        self.scene_manager.grid.draw()

        self.scene_manager.origin.model.write(Matrix44.identity().astype('f4').tobytes())
        self.scene_manager.origin.view.write(m_view.astype('f4').tobytes())
        self.scene_manager.origin.projection.write(m_proj.astype('f4').tobytes())
        self.scene_manager.origin.draw()

        # draw guides
        self.drawSceneObjects(m_view, m_proj)

        # Activate the window screen as the render target
        self.ctx.disable(moderngl.DEPTH_TEST)
        self.screen.clear(0.0)
        self.screen.use()

        # Image from renderer
        self.offscreen2_diffuse.use(location=0)
        self.quad_fs.render()

        # Render offscreen guides buffer over of screen
        self.ctx.enable(moderngl.BLEND)
        self.ctx.blend_equation = moderngl.FUNC_ADD
        self.ctx.blend_func = moderngl.SRC_ALPHA, moderngl.ONE_MINUS_SRC_ALPHA
        self.offscreen_diffuse.use(location=0)
        self.quad_fs.render()
        self.ctx.disable(moderngl.BLEND)

        self.ctx.finish()

        time_now = time.time()
        self.hud_info.fps = 1.0 / (time_now - start_time)

        # render hud surface
        self.renderHUD(0)

    def buildOffscreen(self, width, height):
        buffer_size = (width, height)
        # offscreen render target
        if self.offscreen:
            self.offscreen.release()
            self.offscreen_diffuse.release()
            self.offscreen_normals.release()
            self.offscreen_depth.release()

        self.offscreen_diffuse = self.ctx.texture(buffer_size, 4, dtype='f2') # RGBA color/diffuse layer
        self.offscreen_normals = self.ctx.texture(buffer_size, 4, dtype='f2') # Textures for storing normals (16 bit floats)
        self.offscreen_depth = self.ctx.depth_texture(buffer_size) # Texture for storing depth values

        # create a framebuffer we can render to
        self.offscreen = self.ctx.framebuffer(
            color_attachments=[
                self.offscreen_diffuse,
                self.offscreen_normals
            ],
            depth_attachment=self.offscreen_depth,
        )
        self.offscreen.viewport = (0, 0, width, height)

        if self.offscreen2:
            self.offscreen2.release()
            self.offscreen2_diffuse.release()
            self.offscreen2_depth.release()

        self.offscreen2_diffuse = self.ctx.texture(buffer_size, 4, dtype='f2') # RGBA color/diffuse layer
        self.offscreen2_depth = self.ctx.depth_texture(buffer_size) # Texture for storing depth values

        # create a framebuffer we can render to
        self.offscreen2 = self.ctx.framebuffer(
            color_attachments=[
                self.offscreen2_diffuse,
            ],
            depth_attachment=self.offscreen2_depth,
        )
        self.offscreen2.viewport = (0, 0, width, height)

        # offscreen hud pixmap
        if self.hud_pixmap:
            self.hud_pixmap = None

        if not self.hud_pixmap:
            self.hud_pixmap = QtGui.QPixmap(width, height)
            self.hud_overlay.setPixmap(self.hud_pixmap)

    def resize(self, width, height):
        self._width, self._height = width, height
        self.makeCurrent()
        self.activeCamera.setViewportDimensions(width, height)
        self.ctx.viewport = (0, 0, width, height)

        self.screen = self.ctx.detect_framebuffer(self.defaultFramebufferObject())

        self.renderer.stop()

        self.buildOffscreen(width, height)
        self.hud_overlay.resize(width, height)

        self.renderer.start(width, height) # actually starts/restart rendering

    @property
    def activeCamera(self):
        return self.cameras[self.current_camera]

    def mousePressEvent(self, mouseEvent):
        self.orbit_mode = True
        self.old_mouse_x = mouseEvent.x()
        self.old_mouse_y = mouseEvent.y()
        self.setCursor(QtCore.Qt.ClosedHandCursor)

    def mouseReleaseEvent(self, mouseEvent):
        self.orbit_mode = False
        self.setCursor(QtCore.Qt.OpenHandCursor)

    def mouseMoveEvent(self, mouseEvent):
        if int(mouseEvent.buttons()) != QtCore.Qt.NoButton :
            # user is orbiting camera
            delta_x = mouseEvent.x() - self.old_mouse_x
            delta_y = self.old_mouse_y - mouseEvent.y()

            if int(mouseEvent.buttons()) & QtCore.Qt.LeftButton :
                if mouseEvent.modifiers() == QtCore.Qt.AltModifier :
                    # pan camera
                    self.activeCamera.pan( delta_x, delta_y )
                else:
                    # orbit camera
                    self.activeCamera.orbit(delta_x, delta_y)
            elif int(mouseEvent.buttons()) & QtCore.Qt.RightButton :
                # dolly camera
                self.activeCamera.dolly( 3*(delta_x + delta_y), False )
            elif int(mouseEvent.buttons()) & QtCore.Qt.MidButton :
                # pan camera
                self.activeCamera.pan( delta_x, delta_y )

            self.old_mouse_x = mouseEvent.x()
            self.old_mouse_y = mouseEvent.y()

            self.renderer.updateCamera(self.activeCamera)
            self.update()

    # hou module stuff
    def camera(self) -> ObjNode or None:
        '''
        If the viewport is currently looking through a camera or light (not necessarily locked to it), this returns an object representing the camera/light’s node.
        Returns None if the viewport is not looking through a camera/light.
        '''
        return None


import numpy as np
from robot.robot_scenarios.collaborative_tasks import CollScenario
from robot.robot_core import Robot, Robotworld, Landmark, Hierworld, Hieragent

# set random seed
np.random.seed(2)

class Scenario(CollScenario):
    """Nothing yet here"""

    def make_world(self):
        # define scenario properties
        num_agents = 2
        num_objects = 1
        num_goals = 1
        num_joints = 2
        arm_length = 0.35

        # create world
        world = Hierworld()

        world.agents = [Hieragent()]

        # add robot agents
        world.roboworld.agents = [Robot() for i in range(num_agents)]
        for i, agent in enumerate(world.roboworld.agents):
            agent.name = 'agent %d' % i

        # add objects with landmark properties
        world.roboworld.objects = [Landmark() for i in range(num_objects)]
        for object in world.roboworld.objects:
            object.name = 'object %d' % i

        # add goals with landmark properties
        world.roboworld.goals = [Landmark() for i in range(num_goals)]
        for goal in world.roboworld.goals:
            goal.name = 'goal'

        # add scenario parameters to world
        world.roboworld.num_joints = num_joints
        world.roboworld.arm_length = arm_length

        # reset world
        self.reset_world(world)

        return world

    def reset_world(self, world):
        """Reset the world state """

        # set agent properties
        origins = world.roboworld.robot_position(len(world.roboworld.agents))
        for i, agent in enumerate(world.roboworld.agents):
            agent.color = np.array([0.25,0.25,0.25])
            agent.state.lengths = world.roboworld.arm_length * np.ones(world.roboworld.num_joints)
            agent.state.angles = (2 * np.random.rand(world.roboworld.num_joints) - 1) * np.pi
            agent.state.p_pos = np.array(origins[i][:])

        # set properties for objects
        for object in world.roboworld.objects:
            object.color = np.array([0, 0, 1])
            object.state.p_pos = world.roboworld.object_position(world.roboworld.agents[1].state.p_pos,
                                                                 radius=2*world.roboworld.arm_length)

        # set goal properties
        for goal in world.roboworld.goals:
            goal.state.p_pos = np.array([-1.0, 0.0])
            goal.color = np.array([1, 0, 0])

    def observation(self, agent, world):
        # """Observation function for multi agent collaborative task."""
        #
        # # initialize observation variables
        # state_obs, object_obs, partner_obs, goal_obs = [], [], [], []
        # grasp_obs = [agent.state.grasp]
        # # update agent state observation
        # for i in range(1, world.roboworld.num_joints + 1):
        #     # state_obs += agent.get_joint_pos(i).tolist()
        #     state_obs += agent.local_joint_pos(i).tolist()
        # # update object state observation
        # for i in range(len(world.objects)):
        #     object_obs += np.ndarray.tolist(world.objects[i].state.p_pos - agent.state.p_pos)
        # # update partner observation, if any (centralized observation, not relative)
        # if len(world.agents) > 1:
        #     for partner in world.agents:
        #         # only for partner agents
        #         if agent.name == partner.name: continue
        #         # partner state observations
        #         partner_obs += np.ndarray.tolist(partner.get_joint_pos(world.num_joints) - agent.state.p_pos)
        #         # for i in range(world.num_joints + 1):
        #         #     partner_obs += partner.get_joint_pos(i).tolist()
        #         # add partner grasp observation
        #         partner_obs += [partner.state.grasp]
        # # update goal observation
        # if len(world.goals) > 0:
        #     goal_obs = np.ndarray.tolist(world.goals[0].state.p_pos - agent.state.p_pos)

        return []


# Based on code from astropy
# Licensed under a 3-clause BSD style license

import warnings

import os
import ctypes
from functools import partial

import numpy as np
from numpy.ctypeslib import ndpointer, load_library

from .utils import KernelSizeError, has_even_axis, raise_even_kernel_exception

LIBRARY_PATH = os.path.dirname(__file__)

try:
    _convolve = load_library("_convolve", LIBRARY_PATH)
except Exception:
    raise ImportError("Convolution C extension is missing. Try re-building astropy.")

# The GIL is automatically released by default when calling functions imported
# from libraries loaded by ctypes.cdll.LoadLibrary(<path>)

# Declare prototypes
# Boundary None
_convolveNd_c = _convolve.convolveNd_c
_convolveNd_c.restype = None
_convolveNd_c.argtypes = [ndpointer(ctypes.c_double, flags={"C_CONTIGUOUS", "WRITEABLE"}),  # return array
                          ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"),  # input array
                          ctypes.c_uint,  # N dim
                          # size array for input and result unless
                          # embed_result_within_padded_region is False,
                          # in which case the result array is assumed to be
                          # input.shape - 2*(kernel.shape//2). Note: integer division.
                          ndpointer(ctypes.c_size_t, flags="C_CONTIGUOUS"),
                          ndpointer(ctypes.c_double, flags="C_CONTIGUOUS"),  # kernel array
                          ndpointer(ctypes.c_size_t, flags="C_CONTIGUOUS"),  # size array for kernel
                          ctypes.c_bool,  # nan_interpolate
                          ctypes.c_bool,  # embed_result_within_padded_region
                          ctypes.c_uint]  # n_threads

# Disabling all doctests in this module until a better way of handling warnings
# in doctests can be determined
__doctest_skip__ = ['*']

BOUNDARY_OPTIONS = [None, 'fill', 'wrap', 'extend']

MAX_NORMALIZATION = 100

def _copy_input_if_needed(input, dtype=float, order='C', nan_treatment=None,
                          mask=None, fill_value=None):
    # strip quantity attributes
    if hasattr(input, 'unit'):
        input = input.value
    output = input
    # Copy input
    try:
        # Anything that's masked must be turned into NaNs for the interpolation.
        # This requires copying. A copy is also needed for nan_treatment == 'fill'
        # A copy prevents possible function side-effects of the input array.
        if nan_treatment == 'fill' or np.ma.is_masked(input) or mask is not None:
            if np.ma.is_masked(input):
                # ``np.ma.maskedarray.filled()`` returns a copy, however there
                # is no way to specify the return type or order etc. In addition
                # ``np.nan`` is a ``float`` and there is no conversion to an
                # ``int`` type. Therefore, a pre-fill copy is needed for non
                # ``float`` masked arrays. ``subok=True`` is needed to retain
                # ``np.ma.maskedarray.filled()``. ``copy=False`` allows the fill
                # to act as the copy if type and order are already correct.
                output = np.array(input, dtype=dtype, copy=False, order=order, subok=True)
                output = output.filled(fill_value)
            else:
                # Since we're making a copy, we might as well use `subok=False` to save,
                # what is probably, a negligible amount of memory.
                output = np.array(input, dtype=dtype, copy=True, order=order, subok=False)

            if mask is not None:
                # mask != 0 yields a bool mask for all ints/floats/bool
                output[mask != 0] = fill_value
        else:
            # The call below is synonymous with np.asanyarray(array, ftype=float, order='C')
            # The advantage of `subok=True` is that it won't copy when array is an ndarray subclass. If it
            # is and `subok=False` (default), then it will copy even if `copy=False`. This uses less memory
            # when ndarray subclasses are passed in.
            output = np.array(input, dtype=dtype, copy=False, order=order, subok=True)
    except (TypeError, ValueError) as e:
        raise TypeError('input should be a Numpy array or something '
                        'convertible into a float array', e)
    return output

def convolve(array, kernel, boundary='fill', fill_value=0.,
             nan_treatment='interpolate', normalize_kernel=True, mask=None,
             preserve_nan=False, normalization_zero_tol=1e-8):
    """
    Convolve an array with a kernel.

    This routine differs from `scipy.ndimage.convolve` because
    it includes a special treatment for ``NaN`` values. Rather than
    including ``NaN`` values in the array in the convolution calculation, which
    causes large ``NaN`` holes in the convolved array, ``NaN`` values are
    replaced with interpolated values using the kernel as an interpolation
    function.

    Parameters
    ----------
    array : `numpy.ndarray`
        The array to convolve. This should be a 1, 2, or 3-dimensional array
        or a list or a set of nested lists representing a 1, 2, or
        3-dimensional array.
    kernel : `numpy.ndarray`
        The convolution kernel. The number of dimensions should match those for
        the array, and the dimensions should be odd in all directions.  If a
        masked array, the masked values will be replaced by ``fill_value``.
    boundary : str, optional
        A flag indicating how to handle boundaries:
            * `None`
                Set the ``result`` values to zero where the kernel
                extends beyond the edge of the array.
            * 'fill'
                Set values outside the array boundary to ``fill_value`` (default).
            * 'wrap'
                Periodic boundary that wrap to the other side of ``array``.
            * 'extend'
                Set values outside the array to the nearest ``array``
                value.
    fill_value : float, optional
        The value to use outside the array when using ``boundary='fill'``
    normalize_kernel : bool, optional
        Whether to normalize the kernel to have a sum of one.
    nan_treatment : {'interpolate', 'fill'}
        interpolate will result in renormalization of the kernel at each
        position ignoring (pixels that are NaN in the image) in both the image
        and the kernel.
        'fill' will replace the NaN pixels with a fixed numerical value (default
        zero, see ``fill_value``) prior to convolution
        Note that if the kernel has a sum equal to zero, NaN interpolation
        is not possible and will raise an exception.
    preserve_nan : bool
        After performing convolution, should pixels that were originally NaN
        again become NaN?
    mask : `None` or `numpy.ndarray`
        A "mask" array.  Shape must match ``array``, and anything that is masked
        (i.e., not 0/`False`) will be set to NaN for the convolution.  If
        `None`, no masking will be performed unless ``array`` is a masked array.
        If ``mask`` is not `None` *and* ``array`` is a masked array, a pixel is
        masked of it is masked in either ``mask`` *or* ``array.mask``.
    normalization_zero_tol: float, optional
        The absolute tolerance on whether the kernel is different than zero.
        If the kernel sums to zero to within this precision, it cannot be
        normalized. Default is "1e-8".

    Returns
    -------
    result : `numpy.ndarray`
        An array with the same dimensions and as the input array,
        convolved with kernel.  The data type depends on the input
        array type.  If array is a floating point type, then the
        return array keeps the same data type, otherwise the type
        is ``numpy.float``.

    Notes
    -----
    For masked arrays, masked values are treated as NaNs.  The convolution
    is always done at ``numpy.float`` precision.
    """

    if boundary not in BOUNDARY_OPTIONS:
        raise ValueError("Invalid boundary option: must be one of {}"
                         .format(BOUNDARY_OPTIONS))

    if nan_treatment not in ('interpolate', 'fill'):
        raise ValueError("nan_treatment must be one of 'interpolate','fill'")

    # OpenMP support is disabled at the C src code level, changing this will have
    # no effect.
    n_threads = 1

    # Keep refs to originals
    passed_kernel = kernel
    passed_array = array

    # The C routines all need float type inputs (so, a particular
    # bit size, endianness, etc.).  So we have to convert, which also
    # has the effect of making copies so we don't modify the inputs.
    # After this, the variables we work with will be array_internal, and
    # kernel_internal.  However -- we do want to keep track of what type
    # the input array was so we can cast the result to that at the end
    # if it's a floating point type.  Don't bother with this for lists --
    # just always push those as float.
    # It is always necessary to make a copy of kernel (since it is modified),
    # but, if we just so happen to be lucky enough to have the input array
    # have exactly the desired type, we just alias to array_internal
    # Convert kernel to ndarray if not already

    # Copy or alias array to array_internal
    array_internal = _copy_input_if_needed(passed_array, dtype=float, order='C',
                                           nan_treatment=nan_treatment, mask=mask,
                                           fill_value=np.nan)
    array_dtype = getattr(passed_array, 'dtype', array_internal.dtype)
    # Copy or alias kernel to kernel_internal
    kernel_internal = _copy_input_if_needed(passed_kernel, dtype=float, order='C',
                                            nan_treatment=None, mask=None,
                                            fill_value=fill_value)

    # Make sure kernel has all odd axes
    if has_even_axis(kernel_internal):
        raise_even_kernel_exception()

    # -----------------------------------------------------------------------
    # From this point onwards refer only to ``array_internal`` and
    # ``kernel_internal``.
    # Assume both are base np.ndarrays and NOT subclasses e.g. NOT
    # ``Kernel`` nor ``np.ma.maskedarray`` classes.
    # -----------------------------------------------------------------------

    # Check dimensionality
    if array_internal.ndim == 0:
        raise Exception("cannot convolve 0-dimensional arrays")
    elif array_internal.ndim > 3:
        raise NotImplementedError('convolve only supports 1, 2, and 3-dimensional '
                                  'arrays at this time')
    elif array_internal.ndim != kernel_internal.ndim:
        raise Exception('array and kernel have differing number of '
                        'dimensions.')

    array_shape = np.array(array_internal.shape)
    kernel_shape = np.array(kernel_internal.shape)
    pad_width = kernel_shape//2

    # For boundary=None only the center space is convolved. All array indices within a
    # distance kernel.shape//2 from the edge are completely ignored (zeroed).
    # E.g. (1D list) only the indices len(kernel)//2 : len(array)-len(kernel)//2
    # are convolved. It is therefore not possible to use this method to convolve an
    # array by a kernel that is larger (see note below) than the array - as ALL pixels would be ignored
    # leaving an array of only zeros.
    # Note: For even kernels the correctness condition is array_shape > kernel_shape.
    # For odd kernels it is:
    # array_shape >= kernel_shape OR array_shape > kernel_shape-1 OR array_shape > 2*(kernel_shape//2).
    # Since the latter is equal to the former two for even lengths, the latter condition is complete.
    if boundary is None and not np.all(array_shape > 2*pad_width):
        raise KernelSizeError("for boundary=None all kernel axes must be smaller than array's - "
                              "use boundary in ['fill', 'extend', 'wrap'] instead.")

    # NaN interpolation significantly slows down the C convolution
    # computation. Since nan_treatment = 'interpolate', is the default
    # check whether it is even needed, if not, don't interpolate.
    # NB: np.isnan(array_internal.sum()) is faster than np.isnan(array_internal).any()
    nan_interpolate = (nan_treatment == 'interpolate') and np.isnan(array_internal.sum())

    # Check if kernel is normalizable
    if normalize_kernel or nan_interpolate:
        kernel_sum = kernel_internal.sum()
        kernel_sums_to_zero = np.isclose(kernel_sum, 0, atol=normalization_zero_tol)

        if kernel_sum < 1. / MAX_NORMALIZATION or kernel_sums_to_zero:
            raise ValueError("The kernel can't be normalized, because its sum is "
                             "close to zero. The sum of the given kernel is < {}"
                             .format(1. / MAX_NORMALIZATION))

    # Mark the NaN values so we can replace them later if interpolate_nan is
    # not set
    if preserve_nan or nan_treatment == 'fill':
        initially_nan = np.isnan(array_internal)
        if nan_treatment == 'fill':
            array_internal[initially_nan] = fill_value

    # Avoid any memory allocation within the C code. Allocate output array
    # here and pass through instead.
    result = np.zeros(array_internal.shape, dtype=float, order='C')

    embed_result_within_padded_region = True
    array_to_convolve = array_internal
    if boundary in ('fill', 'extend', 'wrap'):
        embed_result_within_padded_region = False
        if boundary == 'fill':
            # This method is faster than using numpy.pad(..., mode='constant')
            array_to_convolve = np.full(array_shape + 2*pad_width, fill_value=fill_value, dtype=float, order='C')
            # Use bounds [pad_width[0]:array_shape[0]+pad_width[0]] instead of [pad_width[0]:-pad_width[0]]
            # to account for when the kernel has size of 1 making pad_width = 0.
            if array_internal.ndim == 1:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0]] = array_internal
            elif array_internal.ndim == 2:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0],
                                  pad_width[1]:array_shape[1]+pad_width[1]] = array_internal
            else:
                array_to_convolve[pad_width[0]:array_shape[0]+pad_width[0],
                                  pad_width[1]:array_shape[1]+pad_width[1],
                                  pad_width[2]:array_shape[2]+pad_width[2]] = array_internal
        else:
            np_pad_mode_dict = {'fill': 'constant', 'extend': 'edge', 'wrap': 'wrap'}
            np_pad_mode = np_pad_mode_dict[boundary]
            pad_width = kernel_shape // 2

            if array_internal.ndim == 1:
                np_pad_width = (pad_width[0],)
            elif array_internal.ndim == 2:
                np_pad_width = ((pad_width[0],), (pad_width[1],))
            else:
                np_pad_width = ((pad_width[0],), (pad_width[1],), (pad_width[2],))

            array_to_convolve = np.pad(array_internal, pad_width=np_pad_width,
                                       mode=np_pad_mode)

    _convolveNd_c(result, array_to_convolve,
                  array_to_convolve.ndim,
                  np.array(array_to_convolve.shape, dtype=ctypes.c_size_t, order='C'),
                  kernel_internal,
                  np.array(kernel_shape, dtype=ctypes.c_size_t, order='C'),
                  nan_interpolate, embed_result_within_padded_region,
                  n_threads)

    # So far, normalization has only occured for nan_treatment == 'interpolate'
    # because this had to happen within the C extension so as to ignore
    # any NaNs
    if normalize_kernel:
        if not nan_interpolate:
            result /= kernel_sum
    elif nan_interpolate:
        result *= kernel_sum

    if nan_interpolate and not preserve_nan and np.isnan(result.sum()):
        warnings.warn("nan_treatment='interpolate', however, NaN values detected "
                      "post convolution. A contiguous region of NaN values, larger "
                      "than the kernel size, are present in the input array. "
                      "Increase the kernel size to avoid this.")

    if preserve_nan:
        result[initially_nan] = np.nan

    # Convert result to original data type
    if array_dtype.kind == 'f':
        # Try to preserve the input type if it's a floating point type
        # Avoid making another copy if possible
        try:
            return result.astype(array_dtype, copy=False)
        except TypeError:
            return result.astype(array_dtype)
    else:
        return result


import time
import gym, numpy as np
import  pygame, pymunk, logging, math, random
# imports for DQNAgent
import numpy as np
from keras.models import Sequential
from keras.layers import Dense, Dropout, Conv2D, MaxPooling2D, Activation, Flatten
from tensorflow.keras.optimizers import Adam
from keras.callbacks import TensorBoard
import tensorflow as tf
from collections import deque
import time
import random
import os

# Hide GPU from visible devices
tf.config.set_visible_devices([], 'GPU')

DISCOUNT = 0.99
REPLAY_MEMORY_SIZE = 50_000  # How many last steps to keep for model training
MIN_REPLAY_MEMORY_SIZE = 1_000  # Minimum number of steps in a memory to start training
MINIBATCH_SIZE = 64  # How many steps (samples) to use for training
UPDATE_TARGET_EVERY = 5  # Terminal states (end of episodes)
MODEL_NAME = 'BOX'

# Exploration settings
ELIPSON_DECAY = 0.099975
MIN_EPSILON = 0.001

# For stats
ep_rewards = [-200]

# For more repetitive results
random.seed(1)
np.random.seed(1)

# Memory fraction, used mostly when training multiple agents
#gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=MEMORY_FRACTION)
#backend.set_session(tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)))

# Create models folder
if not os.path.isdir('models'):
    os.makedirs('models')

# Own Tensorboard class
class ModifiedTensorBoard(TensorBoard):

    # Overriding init to set initial step and writer (we want one log file for all .fit() calls)
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.step = 1
        self.writer = tf.summary.create_file_writer(self.log_dir)

    # Overriding this method to stop creating default log writer
    def set_model(self, model):
        pass

    # Overrided, saves logs with our step number
    # (otherwise every .fit() will start writing from 0th step)
    def on_epoch_end(self, epoch, logs=None):
        self.update_stats(**logs)

    # Overrided
    # We train for one batch only, no need to save anything at epoch end
    def on_batch_end(self, batch, logs=None):
        pass

    # Overrided, so won't close writer
    def on_train_end(self, _):
        pass

    # Custom method for saving own metrics
    # Creates writer, writes custom metrics and closes writer
    def update_stats(self, **stats):
        self._write_logs(stats, self.step)

# Agent class
class DQNAgent:
    def __init__(self, env):
        self.env = env

        # Main model
        self.model = self.create_model()

        # Target network
        self.target_model = self.create_model()
        self.target_model.set_weights(self.model.get_weights())

        # An array with last n steps for training
        self.replay_memory = deque(maxlen=REPLAY_MEMORY_SIZE)

        # Custom tensorboard object
        self.tensorboard = ModifiedTensorBoard(log_dir="logs/{}-{}".format(MODEL_NAME, int(time.time())))

        # Used to count when to update target network with main network's weights
        self.target_update_counter = 0

    def create_model(self,):
        model = Sequential()

        observation_space = 60000, np.array(self.env.observation).shape[0], np.array(self.env.observation).shape[1], 1
        action_space = self.env.action_space.n

        model.add(Conv2D(32, (3, 3), activation='relu', input_shape=observation_space[1:]))
        model.add(Activation('relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.2))

        model.add(Conv2D(256, (3, 3)))
        model.add(Activation('relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.2))

        model.add(Flatten())  # this converts our 3D feature maps to 1D feature vectors
        model.add(Dense(64))

        model.add(Dense(action_space, activation='linear'))  # ACTION_SPACE_SIZE = how many choices (9)
        model.compile(loss="mse", optimizer=Adam(lr=0.001), metrics=['accuracy'])
        return model

    # Adds step's data to a memory replay array
    # (observation space, action, reward, new observation space, done)
    def update_replay_memory(self, transition):
        self.replay_memory.append(transition)

    # Trains main network every step during episode
    def train(self, terminal_state, step):

        # Start training only if certain number of samples is already saved
        if len(self.replay_memory) < MIN_REPLAY_MEMORY_SIZE:
            return

        # Get a minibatch of random samples from memory replay table
        minibatch = random.sample(self.replay_memory, MINIBATCH_SIZE)

        # Get current states from minibatch, then query NN model for Q values
        current_states = np.array([transition[0] for transition in minibatch])/255
        current_qs_list = self.model.predict(current_states)

        # Get future states from minibatch, then query NN model for Q values
        # When using target network, query it, otherwise main network should be queried
        new_current_states = np.array([transition[3] for transition in minibatch])/255
        future_qs_list = self.target_model.predict(new_current_states)

        X = []
        y = []

        # Now we need to enumerate our batches
        for index, (current_state, action, reward, new_current_state, done) in enumerate(minibatch):

            # If not a terminal state, get new q from future states, otherwise set it to 0
            # almost like with Q Learning, but we use just part of equation here
            if not done:
                max_future_q = np.max(future_qs_list[index])
                new_q = reward + DISCOUNT * max_future_q
            else:
                new_q = reward

            # Update Q value for given state
            current_qs = current_qs_list[index]
            current_qs[action] = new_q

            # And append to our training data
            X.append(current_state)
            y.append(current_qs)

        # Fit on all samples as one batch, log only on terminal state
        self.model.fit(np.array(X)/255, np.array(y), batch_size=MINIBATCH_SIZE, verbose=0, shuffle=False, callbacks=[self.tensorboard] if terminal_state else None)

        # Update target network counter every episode
        if terminal_state:
            self.target_update_counter += 1

        # If counter reaches set value, update target network with weights of main network
        if self.target_update_counter > UPDATE_TARGET_EVERY:
            self.target_model.set_weights(self.model.get_weights())
            self.target_update_counter = 0

    # Queries main network for Q values given current observation space (environment state)
    def get_qs(self, state):
        return self.model.predict(np.array(state).reshape(-1, *state.shape)/255)[0]

class WorldEnvironment(gym.Env):
    def __init__(self, terrain_world, parent):
        self.action_space = gym.spaces.Discrete(15)
        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(84, 84, 1), dtype=np.uint8)

        self.velocity = (0, 0)
        self.position = (0, 0)
        self.terrain_world = terrain_world
        self.parent = parent

        self.inventory = []
        self.scheduled_rewards = []
        self.time_lapsed = 0
        self.health = 100

        self.observation = self.get_observation()
        self.last_state = self.observation

        self.agent = DQNAgent(self)

    def play_sound(self):
        # play self.parent.parent.parent.assets.get("coin.wav")
        sound = self.parent.parent.parent.textures.get("coin.wav")
        pygame.mixer.Sound.play(sound)

    def get_observation(self):
        _ = int(self.position[0] / 32), int(self.position[1] / 32)
        observation = self.parent.parent.get_terrain_matrix(_, fov=25)
        return observation

    def step(self, action):
        self.position = self.parent.body.position
        self.time_lapsed += 1
        reward = 0
        # generate action id from tensor
        if not isinstance(action, int):
            action = action.argmax()
            if action > 8 and action < 13 and random.randint(0, 3) == 0:
                self.play_sound()
        if action == 0:
            self.velocity = (self.velocity[0] + 40, self.velocity[1])
        elif action == 1:
            self.velocity = (self.velocity[0] - 40, self.velocity[1])
        elif action == 2:
            self.velocity = (self.velocity[0], self.velocity[1] - 400)

        elif action == 3:
            self.velocity = (self.velocity[0] + 40, self.velocity[1] - 400)
        elif action == 4:
            self.velocity = (self.velocity[0] - 40, self.velocity[1] - 400)

        elif action == 5:
            # break block above
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            block = self.parent.parent.remove_block(pos)
            if block is not None:
                self.inventory.append(block)
        elif action == 6:
            # break block below
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            block = self.parent.parent.remove_block((pos[0], pos[1] + 1))
            if block is not None:
                self.inventory.append(block)
        elif action == 7:
            # break block left
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            block = self.parent.parent.remove_block((pos[0]-1, pos[1]))
            if block is not None:
                self.inventory.append(block)
        elif action == 8:
            # break block right
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            block = self.parent.parent.remove_block((pos[0]+1, pos[1]))
            if block is not None:
                self.inventory.append(block)

        elif action == 9:
            # place block above
            try:
                pos = int(self.position[0] / 32), int(self.position[1] / 32)
                if len(self.inventory) > 0:
                    block = self.inventory.pop()
                    self.parent.parent.place_block(pos, block)
            except Exception as e:
                pass
        elif action == 10:
            # place block below
            try:
                pos = int(self.position[0] / 32), int(self.position[1] / 32)
                if len(self.inventory) > 0:
                    block = self.inventory.pop()
                    self.parent.parent.place_block((pos[0], pos[1] + 1), block)
            except Exception as e:
                pass
        elif action == 11:
            # place block left
            try:
                pos = int(self.position[0] / 32), int(self.position[1] / 32)
                if len(self.inventory) > 0:
                    block = self.inventory.pop()
                    self.parent.parent.place_block((pos[0]-1, pos[1]), block)
            except Exception as e:
                pass
        elif action == 12:
            # place block right
            try:
                pos = int(self.position[0] / 32), int(self.position[1] / 32)
                if len(self.inventory) > 0:
                    block = self.inventory.pop()
                    self.parent.parent.place_block((pos[0]+1, pos[1]), block)
            except Exception as e:
                pass

        # 13, 14: attack left, right
        elif action == 13:
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            _ = self.parent.parent.attack((pos[0]-1, pos[1]))
            if _ == None:
                reward -= 10
            else:
                print(f'{_} was attacked by {self.parent.name}')
        elif action == 14:
            pos = int(self.position[0] / 32), int(self.position[1] / 32)
            _ = self.parent.parent.attack((pos[0]+1, pos[1]))
            if _ == None:
                reward -= 10
            else:
                print(f'{_} was attacked by {self.parent.name}')

        if self.position[1] > 10000:
            reward += -100
            self.reset()
            print(f"[{self.parent.name}] fell off the world")

        # reward on inventory size
        reward += len(self.inventory) / 10

        if len(self.scheduled_rewards) > 0:
            reward += self.scheduled_rewards.pop(0)

        reward += 100 * 1 / self.time_lapsed

        # get distance to (0, 0)
        distance = math.dist((0, 0), self.position)
        if distance > 1000:
            reward += 1 * distance / 1000
        else:
            reward += -1 * distance / 1000

        # If a block exists at the player's position, the player is suffocated
        _ = (self.position[0] / 32, self.position[1] / 32)
        if self.parent.parent.get_terrain_at(_) != 0:
            reward += -100
            self.health -= 1

        # sort the inventory
        self.inventory = sorted(self.inventory, key=lambda x: x)
        observation = self.get_observation()

        # Health is 0 if the player is dead
        if self.health <= 0:
            reward += -100
            self.reset()
            print(f"[{self.parent.name}] died")

        # give reward for maintaining health
        reward += (self.health - 50) * 0.1

        # train DQNAgent
        self.agent.update_replay_memory((self.last_state, action, reward, observation, True))
        self.last_state = observation
        return observation, reward, False, {}

    def reset(self):
        self.parent.reset()

class Boxlander:
    def __init__(self, name, parent):
        self.parent = parent
        self.name = name
        self.FOV = 10
        self.env = WorldEnvironment(self.parent, self)

        self.frame = 0

        self.body = pymunk.Body(1, 1)
        self.body.position = (0, 0)
        self.body.velocity = (0, 0)
        self.body.angle = 0

        self.shape = pymunk.Circle(self.body, 10)
        self.shape.collision_type = 1
        self.shape.color = (255, 255, 255)
        self.shape.elasticity = 0.95
        self.previous_velocity = (0, 0)

        self.parent.space.add(self.body, self.shape)
        self.epsilon = 0.1

    def reset(self):
        self.body.position = (0, 0)
        self.body.velocity = (0, 0)
        self.body.angle = 0
        self.frame = 0
        self.env.health = 100

    def render(self, window):
        for i in range(len(self.parent.collisions)-1):
            try:
                if self.shape in self.parent.collisions[i].shapes:
                    _ = self.previous_velocity[1] / 5000 * 3
                    if _ > 0.01:
                        self.env.health -= _
                        self.env.scheduled_rewards.append(-_)
                    self.parent.collisions.remove(self.parent.collisions[i])
            except Exception as e:
                pass

        if self.frame % 20 == 0:
            if random.randint(0, 100) < self.epsilon * 100 and self.epsilon > MIN_EPSILON:
                self.env.step(self.env.action_space.sample())
                self.epsilon *= ELIPSON_DECAY
            else:
                _ = self.env.agent.get_qs(self.env.get_observation())
                self.env.step(_)
        # apply force
        # clamp velocity
        self.body.velocity = self.body.velocity[0] + self.env.velocity[0], self.body.velocity[1] + self.env.velocity[1]
        self.env.velocity = (0, 0)
        self.env.position = self.body.position

        self.frame += 1

        # if the current body velocity is not 0, set the previous velocity to the current body velocity
        if self.body.velocity[0] != 0 and self.body.velocity[1] != 0:
            self.previous_velocity = self.body.velocity

        # player color based on health
        if self.env.health > 50:
            c = (255, 255, 255)
        elif self.env.health > 25:
            c = (255, 255, 0)
        else:
            c = (255, abs(int(math.sin(self.frame / 10) * 255)), abs(int(math.sin(self.frame / 10) * 255)))
        pygame.draw.circle(window, c, (int(self.body.position[0] - self.parent.parent.x), int(self.body.position[1]) - self.parent.parent.y), 10)
        # nametag
        font = pygame.font.SysFont("comicsansms", 20)
        text = font.render(self.name, True, (255, 255, 255))
        window.blit(text, (int(self.body.position[0] - self.parent.parent.x) - text.get_width() // 2, int(self.body.position[1]) - self.parent.parent.y + 32 - text.get_height() // 2))
        # draw health bar
        pygame.draw.rect(window, (255, 255, 255), (int(self.body.position[0] - self.parent.parent.x) - 10, int(self.body.position[1]) - self.parent.parent.y - 32, 20, 10))
        color = (0, 255, 0) if self.env.health > 75 else (255, 255, 0) if self.env.health > 50 else (255, 0, 0)
        pygame.draw.rect(window, color, (int(self.body.position[0] - self.parent.parent.x) - 10, int(self.body.position[1]) - self.parent.parent.y - 32, 20 * self.env.health / 100, 10))


import argparse
import json
import os
import random
from test import evaluate
import time

import numpy as np
import torch.distributed as dist
import torch.utils.data.distributed

from apex import amp
from apex.parallel import DistributedDataParallel
from warpctc_pytorch import CTCLoss
from .data.data_loader import (
    AudioDataLoader,
    BucketingSampler,
    DistributedBucketingSampler,
    SpectrogramDataset,
)
from .decoder import GreedyDecoder
from .logger import TensorBoardLogger, VisdomLogger
from .model import DeepSpeech, supported_rnns
from .utils import check_loss, reduce_tensor

parser = argparse.ArgumentParser(description="DeepSpeech training")
parser.add_argument(
    "--train-manifest",
    metavar="DIR",
    help="path to train manifest csv",
    default="data/train_manifest.csv",
)
parser.add_argument(
    "--val-manifest",
    metavar="DIR",
    help="path to validation manifest csv",
    default="data/val_manifest.csv",
)
parser.add_argument(
    "--sample-rate", default=16000, type=int, help="Sample rate"
)
parser.add_argument(
    "--batch-size", default=20, type=int, help="Batch size for training"
)
parser.add_argument(
    "--num-workers",
    default=4,
    type=int,
    help="Number of workers used in data-loading",
)
parser.add_argument(
    "--labels-path",
    default="labels.json",
    help="Contains all characters for transcription",
)
parser.add_argument(
    "--window-size",
    default=0.02,
    type=float,
    help="Window size for spectrogram in seconds",
)
parser.add_argument(
    "--window-stride",
    default=0.01,
    type=float,
    help="Window stride for spectrogram in seconds",
)
parser.add_argument(
    "--window",
    default="hamming",
    help="Window type for spectrogram generation",
)
parser.add_argument(
    "--hidden-size", default=800, type=int, help="Hidden size of RNNs"
)
parser.add_argument(
    "--hidden-layers", default=5, type=int, help="Number of RNN layers"
)
parser.add_argument(
    "--rnn-type",
    default="gru",
    help="Type of the RNN. rnn|gru|lstm are supported",
)
parser.add_argument(
    "--epochs", default=70, type=int, help="Number of training epochs"
)
parser.add_argument(
    "--cuda", dest="cuda", action="store_true", help="Use cuda to train model"
)
parser.add_argument(
    "--lr",
    "--learning-rate",
    default=3e-4,
    type=float,
    help="initial learning rate",
)
parser.add_argument("--momentum", default=0.9, type=float, help="momentum")
parser.add_argument(
    "--max-norm",
    default=400,
    type=int,
    help="Norm cutoff to prevent explosion of gradients",
)
parser.add_argument(
    "--learning-anneal",
    default=1.1,
    type=float,
    help="Annealing applied to learning rate every epoch",
)
parser.add_argument(
    "--silent",
    dest="silent",
    action="store_true",
    help="Turn off progress tracking per iteration",
)
parser.add_argument(
    "--checkpoint",
    dest="checkpoint",
    action="store_true",
    help="Enables checkpoint saving of model",
)
parser.add_argument(
    "--checkpoint-per-batch",
    default=0,
    type=int,
    help="Save checkpoint per batch. 0 means never save",
)
parser.add_argument(
    "--visdom",
    dest="visdom",
    action="store_true",
    help="Turn on visdom graphing",
)
parser.add_argument(
    "--tensorboard",
    dest="tensorboard",
    action="store_true",
    help="Turn on tensorboard graphing",
)
parser.add_argument(
    "--log-dir",
    default="visualize/deepspeech_final",
    help="Location of tensorboard log",
)
parser.add_argument(
    "--log-params",
    dest="log_params",
    action="store_true",
    help="Log parameter values and gradients",
)
parser.add_argument(
    "--id",
    default="Deepspeech training",
    help="Identifier for visdom/tensorboard run",
)
parser.add_argument(
    "--save-folder", default="models/", help="Location to save epoch models"
)
parser.add_argument(
    "--model-path",
    default="models/deepspeech_final.pth",
    help="Location to save best validation model",
)
parser.add_argument(
    "--continue-from", default="", help="Continue from checkpoint model"
)
parser.add_argument(
    "--finetune",
    dest="finetune",
    action="store_true",
    help='Finetune the model from checkpoint "continue_from"',
)
parser.add_argument(
    "--augment",
    dest="augment",
    action="store_true",
    help="Use random tempo and gain perturbations.",
)
parser.add_argument(
    "--noise-dir",
    default=None,
    help="Directory to inject noise into audio. If default, noise Inject not added",
)
parser.add_argument(
    "--noise-prob",
    default=0.4,
    help="Probability of noise being added per sample",
)
parser.add_argument(
    "--noise-min",
    default=0.0,
    help="Minimum noise level to sample from. (1.0 means all noise, not original signal)",
    type=float,
)
parser.add_argument(
    "--noise-max",
    default=0.5,
    help="Maximum noise levels to sample from. Maximum 1.0",
    type=float,
)
parser.add_argument(
    "--no-shuffle",
    dest="no_shuffle",
    action="store_true",
    help="Turn off shuffling and sample from dataset based on sequence length (smallest to largest)",
)
parser.add_argument(
    "--no-sortaGrad",
    dest="no_sorta_grad",
    action="store_true",
    help="Turn off ordering of dataset on sequence length for the first epoch.",
)
parser.add_argument(
    "--no-bidirectional",
    dest="bidirectional",
    action="store_false",
    default=True,
    help="Turn off bi-directional RNNs, introduces lookahead convolution",
)
parser.add_argument(
    "--dist-url",
    default="tcp://127.0.0.1:1550",
    type=str,
    help="url used to set up distributed training",
)
parser.add_argument(
    "--dist-backend", default="nccl", type=str, help="distributed backend"
)
parser.add_argument(
    "--world-size", default=1, type=int, help="number of distributed processes"
)
parser.add_argument(
    "--rank", default=0, type=int, help="The rank of this process"
)
parser.add_argument(
    "--gpu-rank",
    default=None,
    help="If using distributed parallel for multi-gpu, sets the GPU for the process",
)
parser.add_argument(
    "--seed", default=123456, type=int, help="Seed to generators"
)
parser.add_argument("--opt-level", type=str)
parser.add_argument("--keep-batchnorm-fp32", type=str, default=None)
parser.add_argument("--loss-scale", type=str, default=None)

torch.manual_seed(123456)
torch.cuda.manual_seed_all(123456)

def to_np(x):
    return x.cpu().numpy()

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

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

    # Set seeds for determinism
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)
    random.seed(args.seed)

    device = torch.device("cuda" if args.cuda else "cpu")
    args.distributed = args.world_size > 1
    main_proc = True
    device = torch.device("cuda" if args.cuda else "cpu")
    if args.distributed:
        if args.gpu_rank:
            torch.cuda.set_device(int(args.gpu_rank))
        dist.init_process_group(
            backend=args.dist_backend,
            init_method=args.dist_url,
            world_size=args.world_size,
            rank=args.rank,
        )
        main_proc = args.rank == 0  # Only the first proc should save models
    save_folder = args.save_folder
    os.makedirs(save_folder, exist_ok=True)  # Ensure save folder exists

    loss_results, cer_results, wer_results = (
        torch.Tensor(args.epochs),
        torch.Tensor(args.epochs),
        torch.Tensor(args.epochs),
    )
    best_wer = None
    if main_proc and args.visdom:
        visdom_logger = VisdomLogger(args.id, args.epochs)
    if main_proc and args.tensorboard:
        tensorboard_logger = TensorBoardLogger(
            args.id, args.log_dir, args.log_params
        )

    avg_loss, start_epoch, start_iter, optim_state = 0, 0, 0, None
    if args.continue_from:  # Starting from previous model
        print("Loading checkpoint model %s" % args.continue_from)
        package = torch.load(
            args.continue_from, map_location=lambda storage, loc: storage
        )
        model = DeepSpeech.load_model_package(package)
        labels = model.labels
        audio_conf = model.audio_conf
        if not args.finetune:  # Don't want to restart training
            optim_state = package["optim_dict"]
            start_epoch = (
                int(package.get("epoch", 1)) - 1
            )  # Index start at 0 for training
            start_iter = package.get("iteration", None)
            if start_iter is None:
                start_epoch += 1  # We saved model after epoch finished, start at the next epoch.
                start_iter = 0
            else:
                start_iter += 1
            avg_loss = int(package.get("avg_loss", 0))
            loss_results, cer_results, wer_results = (
                package["loss_results"],
                package["cer_results"],
                package["wer_results"],
            )
            best_wer = wer_results[start_epoch]
            if (
                main_proc and args.visdom
            ):  # Add previous scores to visdom graph
                visdom_logger.load_previous_values(start_epoch, package)
            if (
                main_proc and args.tensorboard
            ):  # Previous scores to tensorboard logs
                tensorboard_logger.load_previous_values(start_epoch, package)
    else:
        with open(args.labels_path) as label_file:
            labels = str("".join(json.load(label_file)))

        audio_conf = dict(
            sample_rate=args.sample_rate,
            window_size=args.window_size,
            window_stride=args.window_stride,
            window=args.window,
            noise_dir=args.noise_dir,
            noise_prob=args.noise_prob,
            noise_levels=(args.noise_min, args.noise_max),
        )

        rnn_type = args.rnn_type.lower()
        assert (
            rnn_type in supported_rnns
        ), "rnn_type should be either lstm, rnn or gru"
        model = DeepSpeech(
            rnn_hidden_size=args.hidden_size,
            nb_layers=args.hidden_layers,
            labels=labels,
            rnn_type=supported_rnns[rnn_type],
            audio_conf=audio_conf,
            bidirectional=args.bidirectional,
        )

    decoder = GreedyDecoder(labels)
    train_dataset = SpectrogramDataset(
        audio_conf=audio_conf,
        manifest_filepath=args.train_manifest,
        labels=labels,
        normalize=True,
        augment=args.augment,
    )
    test_dataset = SpectrogramDataset(
        audio_conf=audio_conf,
        manifest_filepath=args.val_manifest,
        labels=labels,
        normalize=True,
        augment=False,
    )
    if not args.distributed:
        train_sampler = BucketingSampler(
            train_dataset, batch_size=args.batch_size
        )
    else:
        train_sampler = DistributedBucketingSampler(
            train_dataset,
            batch_size=args.batch_size,
            num_replicas=args.world_size,
            rank=args.rank,
        )
    train_loader = AudioDataLoader(
        train_dataset,
        num_workers=args.num_workers,
        batch_sampler=train_sampler,
    )
    test_loader = AudioDataLoader(
        test_dataset, batch_size=args.batch_size, num_workers=args.num_workers
    )

    if (not args.no_shuffle and start_epoch != 0) or args.no_sorta_grad:
        print("Shuffling batches for the following epochs")
        train_sampler.shuffle(start_epoch)

    model = model.to(device)
    parameters = model.parameters()
    optimizer = torch.optim.SGD(
        parameters,
        lr=args.lr,
        momentum=args.momentum,
        nesterov=True,
        weight_decay=1e-5,
    )
    if optim_state is not None:
        optimizer.load_state_dict(optim_state)

    model, optimizer = amp.initialize(
        model,
        optimizer,
        opt_level=args.opt_level,
        keep_batchnorm_fp32=args.keep_batchnorm_fp32,
        loss_scale=args.loss_scale,
    )
    if args.distributed:
        model = DistributedDataParallel(model)
    print(model)
    print("Number of parameters: %d" % DeepSpeech.get_param_size(model))

    criterion = CTCLoss()
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()

    for epoch in range(start_epoch, args.epochs):
        model.train()
        end = time.time()
        start_epoch_time = time.time()
        for i, (data) in enumerate(train_loader, start=start_iter):
            if i == len(train_sampler):
                break
            inputs, targets, input_percentages, target_sizes = data
            input_sizes = input_percentages.mul_(int(inputs.size(3))).int()
            # measure data loading time
            data_time.update(time.time() - end)
            inputs = inputs.to(device)

            out, output_sizes = model(inputs, input_sizes)
            out = out.transpose(0, 1)  # TxNxH

            float_out = out.float()  # ensure float32 for loss
            loss = criterion(
                float_out, targets, output_sizes, target_sizes
            ).to(device)
            loss = loss / inputs.size(0)  # average the loss by minibatch

            if args.distributed:
                loss = loss.to(device)
                loss_value = reduce_tensor(loss, args.world_size).item()
            else:
                loss_value = loss.item()

            # Check to ensure valid loss was calculated
            valid_loss, error = check_loss(loss, loss_value)
            if valid_loss:
                optimizer.zero_grad()
                # compute gradient

                with amp.scale_loss(loss, optimizer) as scaled_loss:
                    scaled_loss.backward()
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(), args.max_norm
                )
                optimizer.step()
            else:
                print(error)
                print("Skipping grad update")
                loss_value = 0

            avg_loss += loss_value
            losses.update(loss_value, inputs.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
            if not args.silent:
                print(
                    "Epoch: [{0}][{1}/{2}]\t"
                    "Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t"
                    "Data {data_time.val:.3f} ({data_time.avg:.3f})\t"
                    "Loss {loss.val:.4f} ({loss.avg:.4f})\t".format(
                        (epoch + 1),
                        (i + 1),
                        len(train_sampler),
                        batch_time=batch_time,
                        data_time=data_time,
                        loss=losses,
                    )
                )
            if (
                args.checkpoint_per_batch > 0
                and i > 0
                and (i + 1) % args.checkpoint_per_batch == 0
                and main_proc
            ):
                file_path = "%s/deepspeech_checkpoint_epoch_%d_iter_%d.pth" % (
                    save_folder,
                    epoch + 1,
                    i + 1,
                )
                print("Saving checkpoint model to %s" % file_path)
                torch.save(
                    DeepSpeech.serialize(
                        model,
                        optimizer=optimizer,
                        epoch=epoch,
                        iteration=i,
                        loss_results=loss_results,
                        wer_results=wer_results,
                        cer_results=cer_results,
                        avg_loss=avg_loss,
                    ),
                    file_path,
                )
            del loss, out, float_out

        avg_loss /= len(train_sampler)

        epoch_time = time.time() - start_epoch_time
        print(
            "Training Summary Epoch: [{0}]\t"
            "Time taken (s): {epoch_time:.0f}\t"
            "Average Loss {loss:.3f}\t".format(
                epoch + 1, epoch_time=epoch_time, loss=avg_loss
            )
        )

        start_iter = 0  # Reset start iteration for next epoch
        with torch.no_grad():
            wer, cer, output_data = evaluate(
                test_loader=test_loader,
                device=device,
                model=model,
                decoder=decoder,
                target_decoder=decoder,
            )
        loss_results[epoch] = avg_loss
        wer_results[epoch] = wer
        cer_results[epoch] = cer
        print(
            "Validation Summary Epoch: [{0}]\t"
            "Average WER {wer:.3f}\t"
            "Average CER {cer:.3f}\t".format(epoch + 1, wer=wer, cer=cer)
        )

        values = {
            "loss_results": loss_results,
            "cer_results": cer_results,
            "wer_results": wer_results,
        }
        if args.visdom and main_proc:
            visdom_logger.update(epoch, values)
        if args.tensorboard and main_proc:
            tensorboard_logger.update(epoch, values, model.named_parameters())
            values = {
                "Avg Train Loss": avg_loss,
                "Avg WER": wer,
                "Avg CER": cer,
            }

        if main_proc and args.checkpoint:
            file_path = "%s/deepspeech_%d.pth.tar" % (save_folder, epoch + 1)
            torch.save(
                DeepSpeech.serialize(
                    model,
                    optimizer=optimizer,
                    epoch=epoch,
                    loss_results=loss_results,
                    wer_results=wer_results,
                    cer_results=cer_results,
                ),
                file_path,
            )
        # anneal lr
        for g in optimizer.param_groups:
            g["lr"] = g["lr"] / args.learning_anneal
        print("Learning rate annealed to: {lr:.6f}".format(lr=g["lr"]))

        if main_proc and (best_wer is None or best_wer > wer):
            print(
                "Found better validated model, saving to %s" % args.model_path
            )
            torch.save(
                DeepSpeech.serialize(
                    model,
                    optimizer=optimizer,
                    epoch=epoch,
                    loss_results=loss_results,
                    wer_results=wer_results,
                    cer_results=cer_results,
                ),
                args.model_path,
            )
            best_wer = wer
            avg_loss = 0

        if not args.no_shuffle:
            print("Shuffling batches...")
            train_sampler.shuffle(epoch)


import math
import pickle
import pandas as pd
import numpy as np
import settings

def scoreDecisionTreeClassifier(Speed_sensor, Vibration, Engine_Load, Coolant_Temp, Intake_Pressure, Engine_RPM, Speed_OBD, Intake_Air, Flow_Rate, Throttle_Pos, Voltage, Ambient, Accel, Engine_Oil_Temp, Speed_GPS, GPS_Longitude, GPS_Latitude, GPS_Bearing, GPS_Altitude, Turbo_Boost, Trip_Distance, Litres_Per_km, Accel_Ssor_Total, CO2, Trip_Time, CO_emission, HC_emission, PM_emission, NOx_emission, CO2_emission, Fuel_level, Oil_life, Vibration_alert, VibrationAlert_Total, Vibration_Recent, Turbo_alert, Emission_alert, Fog_control, Engine_control):
    "Output: EM_EVENTPROBABILITY, EM_CLASSIFICATION"

    try:
        _thisModelFit
    except NameError:
        with open(settings.pickle_path + 'DecisionTreeClassifier.pickle', 'rb') as _pFile:
            _thisModelFit = pickle.load(_pFile)

    try:
        if math.isnan(Speed_sensor):
            Speed_sensor = 75.06777416636571
    except TypeError:
        Speed_sensor = 75.06777416636571

    try:
        if math.isnan(Vibration):
            Vibration = 246.95994479089921
    except TypeError:
        Vibration = 246.95994479089921

    try:
        if math.isnan(Engine_Load):
            Engine_Load = 36.06875427988444
    except TypeError:
        Engine_Load = 36.06875427988444

    try:
        if math.isnan(Coolant_Temp):
            Coolant_Temp = 87.7360057782593
    except TypeError:
        Coolant_Temp = 87.7360057782593

    try:
        if math.isnan(Intake_Pressure):
            Intake_Pressure = 119.41133983387505
    except TypeError:
        Intake_Pressure = 119.41133983387505

    try:
        if math.isnan(Engine_RPM):
            Engine_RPM = 1263.593956903816
    except TypeError:
        Engine_RPM = 1263.593956903816

    try:
        if math.isnan(Speed_OBD):
            Speed_OBD = 75.06777416636571
    except TypeError:
        Speed_OBD = 75.06777416636571

    try:
        if math.isnan(Intake_Air):
            Intake_Air = 10.664138678223186
    except TypeError:
        Intake_Air = 10.664138678223186

    try:
        if math.isnan(Flow_Rate):
            Flow_Rate = 23.769685807150594
    except TypeError:
        Flow_Rate = 23.769685807150594

    try:
        if math.isnan(Throttle_Pos):
            Throttle_Pos = 75.25847583170821
    except TypeError:
        Throttle_Pos = 75.25847583170821

    try:
        if math.isnan(Voltage):
            Voltage = 14.171672083784758
    except TypeError:
        Voltage = 14.171672083784758

    try:
        if math.isnan(Ambient):
            Ambient = 7.309738774527507
    except TypeError:
        Ambient = 7.309738774527507

    try:
        if math.isnan(Accel):
            Accel = 22.004593668231607
    except TypeError:
        Accel = 22.004593668231607

    try:
        if math.isnan(Engine_Oil_Temp):
            Engine_Oil_Temp = 84.52690501986277
    except TypeError:
        Engine_Oil_Temp = 84.52690501986277

    try:
        if math.isnan(Speed_GPS):
            Speed_GPS = 77.28586848607259
    except TypeError:
        Speed_GPS = 77.28586848607259

    try:
        if math.isnan(GPS_Longitude):
            GPS_Longitude = 8.569920139091247
    except TypeError:
        GPS_Longitude = 8.569920139091247

    try:
        if math.isnan(GPS_Latitude):
            GPS_Latitude = 48.54705087261828
    except TypeError:
        GPS_Latitude = 48.54705087261828

    try:
        if math.isnan(GPS_Bearing):
            GPS_Bearing = 195.5093294811605
    except TypeError:
        GPS_Bearing = 195.5093294811605

    try:
        if math.isnan(GPS_Altitude):
            GPS_Altitude = 326.22607439508846
    except TypeError:
        GPS_Altitude = 326.22607439508846

    try:
        if math.isnan(Turbo_Boost):
            Turbo_Boost = 3.2505074667148186
    except TypeError:
        Turbo_Boost = 3.2505074667148186

    try:
        if math.isnan(Trip_Distance):
            Trip_Distance = 230.49551875887803
    except TypeError:
        Trip_Distance = 230.49551875887803

    try:
        if math.isnan(Litres_Per_km):
            Litres_Per_km = 4.316090813719474
    except TypeError:
        Litres_Per_km = 4.316090813719474

    try:
        if math.isnan(Accel_Ssor_Total):
            Accel_Ssor_Total = -0.0031587095220898037
    except TypeError:
        Accel_Ssor_Total = -0.0031587095220898037

    try:
        if math.isnan(CO2):
            CO2 = 93.42142250399041
    except TypeError:
        CO2 = 93.42142250399041

    try:
        if math.isnan(Trip_Time):
            Trip_Time = 7512.937041049717
    except TypeError:
        Trip_Time = 7512.937041049717

    try:
        if math.isnan(CO_emission):
            CO_emission = 0.0
    except TypeError:
        CO_emission = 0.0

    try:
        if math.isnan(HC_emission):
            HC_emission = 0.0
    except TypeError:
        HC_emission = 0.0

    try:
        if math.isnan(PM_emission):
            PM_emission = 0.0
    except TypeError:
        PM_emission = 0.0

    try:
        if math.isnan(NOx_emission):
            NOx_emission = 0.0
    except TypeError:
        NOx_emission = 0.0

    try:
        if math.isnan(CO2_emission):
            CO2_emission = 0.0
    except TypeError:
        CO2_emission = 0.0

    try:
        if math.isnan(Fuel_level):
            Fuel_level = 0.0
    except TypeError:
        Fuel_level = 0.0

    try:
        if math.isnan(Oil_life):
            Oil_life = 0.0
    except TypeError:
        Oil_life = 0.0

    try:
        if math.isnan(Vibration_alert):
            Vibration_alert = 1.0
    except TypeError:
        Vibration_alert = 1.0

    try:
        if math.isnan(VibrationAlert_Total):
            VibrationAlert_Total = 123.0
    except TypeError:
        VibrationAlert_Total = 123.0

    try:
        if math.isnan(Vibration_Recent):
            Vibration_Recent = 12.0
    except TypeError:
        Vibration_Recent = 12.0

    try:
        if math.isnan(Turbo_alert):
            Turbo_alert = 1.0
    except TypeError:
        Turbo_alert = 1.0

    try:
        if math.isnan(Emission_alert):
            Emission_alert = 1.0
    except TypeError:
        Emission_alert = 1.0

    try:
        if math.isnan(Fog_control):
            Fog_control = 1.0
    except TypeError:
        Fog_control = 1.0

    try:
        if math.isnan(Engine_control):
            Engine_control = 1.0
    except TypeError:
        Engine_control = 1.0

    inputArray = pd.DataFrame([[1.0, Speed_sensor, Vibration, Engine_Load, Coolant_Temp, Intake_Pressure, Engine_RPM, Speed_OBD, Intake_Air, Flow_Rate, Throttle_Pos, Voltage, Ambient, Accel, Engine_Oil_Temp, Speed_GPS, GPS_Longitude, GPS_Latitude, GPS_Bearing, GPS_Altitude, Turbo_Boost, Trip_Distance, Litres_Per_km, Accel_Ssor_Total, CO2, Trip_Time, CO_emission, HC_emission, PM_emission, NOx_emission, CO2_emission, Fuel_level, Oil_life, Vibration_alert, VibrationAlert_Total, Vibration_Recent, Turbo_alert, Emission_alert, Fog_control, Engine_control]],
                              columns = ['const', 'Speed_sensor', 'Vibration', 'Engine_Load', 'Coolant_Temp', 'Intake_Pressure', 'Engine_RPM', 'Speed_OBD', 'Intake_Air', 'Flow_Rate', 'Throttle_Pos', 'Voltage', 'Ambient', 'Accel', 'Engine_Oil_Temp', 'Speed_GPS', 'GPS_Longitude', 'GPS_Latitude', 'GPS_Bearing', 'GPS_Altitude', 'Turbo_Boost', 'Trip_Distance', 'Litres_Per_km', 'Accel_Ssor_Total', 'CO2', 'Trip_Time', 'CO_emission', 'HC_emission', 'PM_emission', 'NOx_emission', 'CO2_emission', 'Fuel_level', 'Oil_life', 'Vibration_alert', 'VibrationAlert_Total', 'Vibration_Recent', 'Turbo_alert', 'Emission_alert', 'Fog_control', 'Engine_control'],
                              dtype = float)

    prediction = _thisModelFit.predict(inputArray)

    EM_EVENTPROBABILITY = float(prediction)

    if (EM_EVENTPROBABILITY >= 0.2269170579029734):
        EM_CLASSIFICATION = '1'
    else:
        EM_CLASSIFICATION = '0'

    return(EM_EVENTPROBABILITY, EM_CLASSIFICATION)


# -*- coding: utf-8 -*-

import numpy as np
import tensorflow as tf
import yolo_v3
import yolo_v3_tiny
# from PIL import Image, ImageDraw
#import yolov3_tiny_tflite
import os
import yolo_v3_tiny
from utils import load_weights, load_coco_names, detections_boxes, detections_boxes, gpu_nms
from tensorflow.python.tools import freeze_graph

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string(
    'class_names', "./sur/sur0316/sur.names", 'File with class names')
tf.app.flags.DEFINE_string(
    'weights_file', './sur/sur0316/yolov3_tiny_sur_best.weights', 'Binary file with detector weights')
tf.app.flags.DEFINE_string(
    'ckpt_file', './sur/sur0316/yolov3_tiny_sur.ckpt', 'Binary file with detector weights')
tf.app.flags.DEFINE_string(
    'data_format', 'NCHW', 'Data format: NCHW (gpu only) / NHWC')
tf.app.flags.DEFINE_string(
    'output_graph', './sur/sur0316/yolov3_tiny_sur.pb', 'Frozen tensorflow protobuf model output path')
tf.app.flags.DEFINE_string(
    'output_dir', './sur/sur0316/', 'Frozen tensorflow protobuf model output path')
tf.app.flags.DEFINE_string(
    'output_file', 'yolov3_tiny_sur.pb', 'Frozen tensorflow protobuf model output path')

tf.app.flags.DEFINE_bool(
    'tiny', True, 'Use tiny version of YOLOv3')
tf.app.flags.DEFINE_bool(
    'spp', False, 'Use SPP version of YOLOv3')
tf.app.flags.DEFINE_integer(
    'size', 608, 'Image size')

def build_detection_graph():
    input_data = tf.placeholder(dtype=tf.uint8, shape=[FLAGS.size, FLAGS.size, 3], name='input_data')
    input_data = tf.expand_dims(input_data, 0)
    input_data = tf.cast(input_data, tf.float32)
    input_data = input_data / 255.
    if FLAGS.tiny:
        model = yolo_v3_tiny.yolo_v3_tiny
    elif FLAGS.spp:
        model = yolo_v3.yolo_v3_spp
    else:
        model = yolo_v3.yolo_v3
    classes = load_coco_names(FLAGS.class_names)

   # yolo_model = model(cfgs.class_num, cfgs.anchors)
    with tf.variable_scope('detector'):
        detections = model(input_data, len(classes), data_format=FLAGS.data_format)
        print(detections.get_shape().as_list())
    boxes, pred_confs, pred_probs = tf.split(detections, [4, 1, len(classes)], axis=-1)
    center_x, center_y, width, height = tf.split(boxes, [1, 1, 1, 1], axis=-1)
    x_min = center_x - width / 2
    y_min = center_y - height / 2
    x_max = center_x + width / 2
    y_max = center_y + height / 2
    pred_boxes = tf.concat([x_min, y_min, x_max, y_max], axis=-1)

    pred_scores = pred_confs * pred_probs
    boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, len(classes), max_boxes=20, score_thresh=0.3,
                                    nms_thresh=0.4)

    boxes = tf.identity(boxes, name='boxes')
    scores = tf.identity(scores, name='scores')
    labels = tf.identity(labels, name='labels')

    return boxes, scores, labels

def export_frozenPB():

    tf.reset_default_graph()
    dets = build_detection_graph()

    saver = tf.train.Saver()
    output_node_names = 'boxes,scores,labels'

    with tf.Session() as sess:
        print("we have restred the weights from =====>>\n", FLAGS.ckpt_file)
        saver.restore(sess, FLAGS.ckpt_file)

        tf.train.write_graph(sess.graph_def, FLAGS.output_dir, FLAGS.output_file)
        freeze_graph.freeze_graph(input_graph=FLAGS.output_graph,
                                  input_saver='',
                                  input_binary=False,
                                  input_checkpoint=FLAGS.ckpt_file,
                                  output_node_names=output_node_names,
                                  restore_op_name="save/restore_all",
                                  filename_tensor_name='save/Const:0',
                                  output_graph=FLAGS.output_graph,
                                  clear_devices=False,
                                  initializer_nodes='')

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    export_frozenPB()