import logging
logger = logging.getLogger(__name__)

from gi.repository import Gdk
from gi.repository import GdkPixbuf
from gi.repository import Gtk
import csv

from gettext import gettext as _

import advene.core.config as config

from advene.model.annotation import Annotation
from advene.gui.views import AdhocView

import advene.gui.edit.elements
import advene.gui.popup

import advene.util.helper as helper
from advene.gui.util import dialog, png_to_pixbuf, contextual_drag_begin, contextual_drag_end
from advene.gui.util.completer import Completer

COLUMN_ELEMENT=0
COLUMN_CONTENT=1
COLUMN_TYPE=2
COLUMN_ID=3
COLUMN_BEGIN=4
COLUMN_END=5
COLUMN_DURATION=6
COLUMN_BEGIN_FORMATTED=7
COLUMN_END_FORMATTED=8
COLUMN_PIXBUF=9
COLUMN_COLOR=10
COLUMN_SOURCE_PACKAGE=11
COLUMN_CUSTOM_FIRST=12

name="Element tabular view plugin"

def register(controller):
    controller.register_viewclass(AnnotationTable)
    controller.register_viewclass(GenericTable)

class AnnotationTable(AdhocView):
    view_name = _("Annotation table view")
    view_id = 'table'
    tooltip=_("Display annotations in a table")

    def __init__(self, controller=None, parameters=None, custom_data=None, elements=None, source=None):
        """We can initialize the table using either a list of elements, or  a TALES source expression.

        If both are specified, the list of elements takes precedence.
        """
        super(AnnotationTable, self).__init__(controller=controller)
        self.registered_rules = []
        self.close_on_package_load = False
        self.contextual_actions = (
            (_("Export as CSV"), self.csv_export),
            )
        self.controller=controller

        opt, arg = self.load_parameters(parameters)
        self.options.update(opt)
        a=dict(arg)
        if source is None and 'source' in a:
            source=a['source']

        self.source = source
        if elements is None and source:
            elements = self.get_elements_from_source(source)

        self.elements = elements
        self.options={ 'confirm-time-update': True }

        self.mouseover_annotation = None
        self.last_edited_path = None

        self.model = self.build_model(elements, custom_data)
        self.widget = self.build_widget(custom_data)

        self.registered_rules.append( controller.event_handler.internal_rule (event="SnapshotUpdate",
                                                                              method=self.update_snapshot)
                                      )
        def unregister(*p):
            for r in self.registered_rules:
                self.controller.event_handler.remove_rule(r, type_="internal")
        self.widget.connect('destroy', unregister)

    def get_save_arguments(self):
        if self.source is not None:
            arguments = [ ('source', self.source) ]
        else:
            arguments = []
        return self.options, arguments

    def update_annotation(self, annotation=None, event=None):
        if self.source:
            # Re-evaluate source parameter, in case the annotation was
            # created.
            elements = self.get_elements_from_source(self.source)
        else:
            elements = self.elements

        if elements is not None and (annotation in elements or event.endswith('Delete')):
            self.set_elements(elements)

    def update_snapshot(self, context, parameters):
        pos = int(context.globals['position'])
        media = context.globals['media']
        eps = self.controller.package.imagecache.precision
        for r in self.widget.treeview.get_model():
            if (r[COLUMN_ELEMENT].media == media
                and abs(r[COLUMN_BEGIN] - pos) <= eps):
                # Update pixbuf
                r[COLUMN_PIXBUF] = png_to_pixbuf(self.controller.package.imagecache[pos],
                                                 height=32)

    def get_elements(self):
        """Return the list of elements in their displayed order.

        If a selection is active, return only selected elements.
        """
        selection = self.widget.treeview.get_selection()
        r=selection.count_selected_rows()
        if r == 0 or r == 1:
            selection.select_all()
        store, paths=selection.get_selected_rows()
        return [ store.get_value (store.get_iter(p), COLUMN_ELEMENT) for p in paths ]

    def build_model(self, elements, custom_data=None):
        """Build the ListStore containing the data.

        See set_element docstring for the custom_data method explanation.
        """
        if custom_data is not None:
            custom = custom_data
        else:
            def custom(a):
                return tuple()
        args = (object, str, str, str, int, int, str, str, str, GdkPixbuf.Pixbuf, str, str) + custom(None)
        l=Gtk.ListStore(*args)
        if not elements:
            return l
        for a in elements:
            if isinstance(a, Annotation):
                l.append( (a,
                           self.controller.get_title(a),
                           self.controller.get_title(a.type),
                           a.id,
                           a.fragment.begin,
                           a.fragment.end,
                           helper.format_time(a.fragment.duration),
                           helper.format_time(a.fragment.begin),
                           helper.format_time(a.fragment.end),
                           png_to_pixbuf(self.controller.get_snapshot(annotation=a),
                                         height=32),
                           self.controller.get_element_color(a),
                           a.ownerPackage.getTitle()
                           ) + custom(a),
                          )
        return l

    def set_elements(self, elements, custom_data=None):
        """Use a new set of elements.

        If custom_data is not None, then it is a function returning
        tuples, that can be used to defined additional model columns.

        When called with None as parameter, it must return a tuple
        with the additional column types. It will be appended at the
        end of the ListStore, in columns COLUMN_CUSTOM_FIRST,
        COLUMN_CUSTOM_FIRST+1, etc.

        When called with an annotation as parameter, it must return a
        tuple with the appropriate values for the annotation in the
        custom columns.
        """
        if elements is None:
            elements = []
        model=self.build_model(elements, custom_data)
        self.widget.treeview.set_model(model)
        self.model = model
        self.elements=elements
        if self.last_edited_path is not None:
            # We just edited an annotation. This update must come from
            # it, so let us try to set the cursor position at the next element.
            path = self.last_edited_path.next()
            try:
                self.model.get_iter(path)
            except (ValueError, TypeError):
                path = self.last_edited_path
            self.widget.treeview.set_cursor(path,
                                            self.columns['id'],
                                            True)
            self.last_edited_path = None

    def motion_notify_event_cb(self, tv, event):
        if not event.get_window() is tv.get_bin_window():
            return False
        if event.is_hint:
            pointer = event.get_window().get_pointer()
            x = pointer.x
            y = pointer.y
        else:
            x = int(event.x)
            y = int(event.y)
        t = tv.get_path_at_pos(x, y)
        if t is not None:
            path, col, cx, cy = t
            it = self.model.get_iter(path)
            ann = self.model.get_value(it,
                                       COLUMN_ELEMENT)
            if self.mouseover_annotation != ann:
                # Update
                if self.mouseover_annotation is not None:
                    self.controller.notify('BookmarkUnhighlight', timestamp=self.mouseover_annotation.fragment.begin, media=self.mouseover_annotation.media, immediate=True)
                self.controller.notify('BookmarkHighlight', timestamp=ann.fragment.begin, media=ann.media, immediate=True)
                self.mouseover_annotation = ann
        return False

    def leave_notify_event_cb(self, tv, event):
        if self.mouseover_annotation is not None:
            self.controller.notify('BookmarkUnhighlight', timestamp=self.mouseover_annotation.fragment.begin, media=self.mouseover_annotation.media, immediate=True)
            self.mouseover_annotation = None
        return False

    def build_widget(self, custom_data=None):
        tree_view = Gtk.TreeView(self.model)

        select = tree_view.get_selection()
        select.set_mode(Gtk.SelectionMode.MULTIPLE)

        tree_view.connect('button-press-event', self.tree_view_button_cb)
        tree_view.connect('key-press-event', self.tree_view_key_cb)
        tree_view.connect('row-activated', self.row_activated_cb)
        tree_view.connect('motion-notify-event', self.motion_notify_event_cb)
        tree_view.connect('leave-notify-event', self.leave_notify_event_cb)
        # Deactivate starting search by simply typing. Users have to use the search shortcut (Control-F)
        tree_view.set_enable_search(False)

        def search_content(model, column, key, it):
            if key in model.get_value(it, COLUMN_CONTENT):
                return False
            return True

        tree_view.set_search_equal_func(search_content)

        columns={}

        columns['snapshot']=Gtk.TreeViewColumn(_("Snapshot"), Gtk.CellRendererPixbuf(), pixbuf=COLUMN_PIXBUF)
        columns['snapshot'].set_reorderable(True)
        tree_view.append_column(columns['snapshot'])

        def cell_edited(cell, path_string, text):
            it = self.model.get_iter_from_string(path_string)
            if not it:
                return
            a = self.model.get_value (it, COLUMN_ELEMENT)
            new_content = helper.title2content(text,
                                               a.content,
                                               a.type.getMetaData(config.data.namespace, "representation"))
            if new_content is None:
                self.log(_("Cannot update the annotation, its representation is too complex"))
            elif a.content.data != new_content:
                self.last_edited_path = Gtk.TreePath.new_from_string(path_string)
                self.controller.notify('EditSessionStart', element=a)
                a.content.data = new_content
                self.controller.notify('AnnotationEditEnd', annotation=a)
                self.controller.notify('EditSessionEnd', element=a)
            return True

        def validate_entry_on_focus_out(widget, event, cell, path):
            cell.emit("edited", path, widget.get_text())
            return True

        def entry_editing_started(cell, editable, path):
            if isinstance(editable, Gtk.Entry):
                it = self.model.get_iter_from_string(path)
                if not it:
                    return
                el = self.model.get_value(it, COLUMN_ELEMENT)
                editable._completer = Completer(textview=editable,
                                                controller=self.controller,
                                                element=el,
                                                indexer=el.rootPackage._indexer)

        custom_cols = []
        if custom_data is not None:
            custom_cols = [ ('custom%d' % i,
                             "Custom %d" % i,
                             COLUMN_CUSTOM_FIRST + i)
                            for i in range(len(custom_data(None))) ]
        for (name, label, col) in (
                ('content', _("Content"), COLUMN_CONTENT),
                *custom_cols,
                ('type', _("Type"), COLUMN_TYPE),
                ('begin', _("Begin"), COLUMN_BEGIN_FORMATTED),
                ('end', _("End"), COLUMN_END_FORMATTED),
                ('duration', _("Duration"), COLUMN_DURATION),
                ('id', _("Id"), COLUMN_ID),
                ('package', _("Package"), COLUMN_SOURCE_PACKAGE)
        ):
            renderer = Gtk.CellRendererText()
            columns[name]=Gtk.TreeViewColumn(label, renderer, text=col)
            if name == 'content':
                renderer.connect('editing-started', entry_editing_started)
                renderer.connect('edited', cell_edited)
                renderer.props.editable = True

            columns[name].set_reorderable(True)
            columns[name].set_sort_column_id(col)
            tree_view.append_column(columns[name])

        # Column-specific settings
        columns['begin'].set_sort_column_id(COLUMN_BEGIN)
        columns['end'].set_sort_column_id(COLUMN_END)
        self.model.set_sort_column_id(COLUMN_BEGIN, Gtk.SortType.ASCENDING)
        columns['type'].add_attribute(columns['type'].get_cells()[0],
                                      'cell-background',
                                      COLUMN_COLOR)

        # Resizable columns: content, type
        for name in ('content', 'type', 'snapshot', *[ c[0] for c in custom_cols ]):
            columns[name].set_sizing(Gtk.TreeViewColumnSizing.FIXED)
            columns[name].set_resizable(True)
            columns[name].set_min_width(40)
        columns['content'].set_expand(True)
        columns['content'].set_max_width(800)

        # Allow user classes to tweak behaviour
        self.columns = columns

        # Drag and drop for annotations
        tree_view.drag_source_set(Gdk.ModifierType.BUTTON1_MASK,
                                  config.data.get_target_types('annotation',
                                                               'text-plain',
                                                               'TEXT',
                                                               'STRING'),
                                  Gdk.DragAction.LINK | Gdk.DragAction.COPY | Gdk.DragAction.MOVE)

        def get_element():
            selection = tree_view.get_selection ()
            if not selection:
                return None
            store, paths=selection.get_selected_rows()
            l=[ store.get_value (store.get_iter(p), COLUMN_ELEMENT) for p in paths ]
            if not l:
                return None
            elif len(l) == 1:
                return l[0]
            else:
                return l
        tree_view.connect('drag-begin', contextual_drag_begin, get_element, self.controller)
        tree_view.connect('drag-end', contextual_drag_end)

        tree_view.connect('drag-data-get', self.drag_data_get_cb)

        # The widget can receive drops
        def drag_received_cb(widget, context, x, y, selection, targetType, time):
            """Handle the drop of an annotation type.
            """
            if Gtk.drag_get_source_widget(context).is_ancestor(self.widget):
                # Ignore drops from our own widget
                return False

            if targetType == config.data.target_type['annotation']:
                sources=[ self.controller.package.annotations.get(uri) for uri in str(selection.get_data(), 'utf8').split('\n') ]
                if sources:
                    self.set_elements(sources)
                return True
            elif targetType == config.data.target_type['annotation-type']:
                sources=[ self.controller.package.annotationTypes.get(uri) for uri in str(selection.get_data(), 'utf8').split('\n') ]
                if sources:
                    self.set_elements(sources[0].annotations)
                return True
            return False

        tree_view.connect('drag-data-received', drag_received_cb)
        tree_view.drag_dest_set(Gtk.DestDefaults.MOTION |
                                Gtk.DestDefaults.HIGHLIGHT |
                                Gtk.DestDefaults.ALL,
                                config.data.get_target_types('annotation', 'annotation-type'),
                                Gdk.DragAction.COPY | Gdk.DragAction.LINK | Gdk.DragAction.MOVE)

        sw = Gtk.ScrolledWindow()
        sw.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        sw.add(tree_view)

        sw.treeview = tree_view

        return sw

    def drag_data_get_cb(self, treeview, context, selection, targetType, timestamp):
        model, paths = treeview.get_selection().get_selected_rows()

        els=[ model[p][COLUMN_ELEMENT] for p in paths ]

        if targetType == config.data.target_type['annotation']:
            selection.set(selection.get_target(), 8, "\n".join( e.uri
                                                                for e in els
                                                                if isinstance(e, Annotation) ).encode('utf8'))
            return True
        elif targetType == config.data.target_type['annotation-type']:
            selection.set(selection.get_target(), 8, "\n".join( e.type.uri
                                                                for e in els
                                                                if isinstance(e, Annotation) ).encode('utf8'))
            return True
        elif targetType in (config.data.target_type['text-plain'],
                            config.data.target_type['TEXT'],
                            config.data.target_type['STRING']):
            selection.set(selection.get_target(), 8, "\n".join(e.content.data
                                                               for e in els
                                                               if isinstance(e, Annotation) ).encode('utf8'))
        else:
            logger.warning("Unknown target type for drag: %d", targetType)
        return True

    def get_selected_nodes(self, with_path=False):
        """Return the currently selected nodes.
        """
        selection = self.widget.treeview.get_selection()
        store, paths = selection.get_selected_rows()
        if with_path:
            return [ (store.get_value(store.get_iter(p), COLUMN_ELEMENT), p) for p in paths ]
        else:
            return [ store.get_value(store.get_iter(p), COLUMN_ELEMENT) for p in paths ]

    def get_selected_node(self, with_path=False):
        """Return the currently selected node.

        None if no node is selected or multiple nodes are selected.
        """
        nodes = self.get_selected_nodes(with_path)
        if len(nodes) != 1:
            return None
        else:
            return nodes[0]

    def debug_cb (self, *p, **kw):
        logger.debug("Parameters: %s\nkw: %s", str(p), str(kw))

    def csv_export(self, name=None):
        if name is None:
            name=dialog.get_filename(title=_("Export data to file..."),
                                     default_file="advene_data.csv",
                                     action=Gtk.FileChooserAction.SAVE,
                                     button=Gtk.STOCK_SAVE)
        if name is None:
            return True
        try:
            f=open(name, 'w', encoding='utf-8')
        except IOError as e:
            dialog.message_dialog(label=_("Error while exporting data to %(filename)s: %(error)s"
                                          % {
                                              'filename': name,
                                              'error': str(e),
                                          }), icon=Gtk.MessageType.ERROR)
        w=csv.writer(f)
        tv=self.widget.treeview
        store, paths=tv.get_selection().get_selected_rows()
        source=[ store.get_iter(p) for p in paths ]
        if not source:
            source=tv.get_model()
        w.writerow( (_("id"), _("type"), _("begin"), _("end"), _("content")) )
        for r in source:
            w.writerow( (r[COLUMN_ID], str(r[COLUMN_TYPE]).encode('utf-8'), r[COLUMN_BEGIN], r[COLUMN_END], str(r[COLUMN_ELEMENT].content.data).encode('utf-8') ) )
        f.close()
        self.log(_("Data exported to %s") % name)

    def row_activated_cb(self, widget, path, view_column):
        """Edit the element on Return or double click
        """
        widget.set_cursor(path, self.columns['content'], True)
        return False

    def set_time(self, attr):
        """Sets the time of the current annotation to the current player time.
        """
        an, an_path = self.get_selected_node(with_path=True)
        if an is None:
            return
        current_time = self.controller.player.current_position_value
        confirm = True
        if self.options['confirm-time-update']:
            confirm = dialog.message_dialog(_("Set %(attr)s time to %(time)s") % {
                'attr': _(attr),
                'time': helper.format_time(current_time)
                }, icon=Gtk.MessageType.QUESTION)
        if confirm:
            self.last_edited_path = an_path
            self.controller.notify('EditSessionStart', element=an, immediate=True)
            setattr(an.fragment, attr, current_time)
            self.controller.notify("AnnotationEditEnd", annotation=an)
            self.controller.notify('EditSessionEnd', element=an)

    def tree_view_key_cb(self, widget=None, event=None):
        if event.keyval == Gdk.KEY_space:
            # Space: goto annotation
            ann = self.get_selected_node ()
            if ann is not None:
                self.controller.update_status (status="seek", position=ann.fragment.begin)
                self.controller.gui.set_current_annotation(ann)
                return True
        elif event.keyval == Gdk.KEY_Return and event.get_state() & Gdk.ModifierType.CONTROL_MASK:
            # Control-return: edit annotation
            ann = self.get_selected_node ()
            if ann is not None:
                self.controller.gui.edit_element(ann)
                return True
        elif event.keyval == Gdk.KEY_less and event.get_state() & Gdk.ModifierType.CONTROL_MASK:
            # Control-< : set begin time
            self.set_time('begin')
            return True
        elif event.keyval == Gdk.KEY_greater and event.get_state() & Gdk.ModifierType.CONTROL_MASK:
            # Control-> : set end time
            self.set_time('end')
            return True
        elif event.keyval in (Gdk.KEY_1,
                              Gdk.KEY_2,
                              Gdk.KEY_3,
                              Gdk.KEY_4,
                              Gdk.KEY_5,
                              Gdk.KEY_6,
                              Gdk.KEY_7,
                              Gdk.KEY_8,
                              Gdk.KEY_9):
            ann, path = self.get_selected_node(with_path=True)
            if  helper.get_type_predefined_completions(ann.type):
                # Shortcut for 1 key edition
                index = event.keyval - 48 - 1
                ret = self.controller.quick_completion_fill_annotation(ann, index)
                if ret:
                    self.last_edited_path = path
                # Gtk.TreePath.new_from_string(path_string)
                return ret
            return False

        return False

    def tree_view_button_cb(self, widget=None, event=None):
        if not event.get_window() is widget.get_bin_window():
            return False

        retval = False
        button = event.button
        x = int(event.x)
        y = int(event.y)

        model = self.model
        t = widget.get_path_at_pos(x, y)
        if t is not None:
            path, col, cx, cy = t
            it = model.get_iter(path)
            node = model.get_value(it,
                                   COLUMN_ELEMENT)
            widget.get_selection().select_path (path)
            if button == 3:
                menu = advene.gui.popup.Menu(node, controller=self.controller)
                menu.popup()
                retval = True
            elif button == 1 and col.get_title() == _("Snapshot"):
                # Click on snapshot -> play
                self.controller.update_status("seek", node.fragment.begin)
                # Allow further processing
                retval = False
        return retval

class GenericTable(AdhocView):
    view_name = _("Generic table view")
    view_id = 'generictable'
    tooltip=_("Display Advene elements in a table.")

    def __init__(self, controller=None, parameters=None, elements=None, source=None):
        super(GenericTable, self).__init__(controller=controller)
        self.close_on_package_load = False
        self.contextual_actions = ()
        self.controller=controller
        self.elements=elements
        self.source = source
        self.options = { }

        opt, arg = self.load_parameters(parameters)
        self.options.update(opt)
        a=dict(arg)
        if source is None and 'source' in a:
            source=a['source']

        if elements is None and source:
            c=self.controller.build_context()
            try:
                elements = c.evaluateValue(source)
                self.source = source
            except Exception as e:
                self.log(_("Error in source evaluation %(source)s: %(error)s") % {
                    'source': self.source,
                    'error': str(e) })
                elements = []

        self.model=self.build_model(elements)
        self.widget = self.build_widget()

    def get_save_arguments(self):
        if self.source is not None:
            arguments = [ ('source', self.source) ]
        else:
            arguments = []
        return self.options, arguments

    def get_elements(self):
        """Return the list of elements in their displayed order.

        If a selection is active, return only selected elements.
        """
        selection = self.widget.treeview.get_selection ()
        r=selection.count_selected_rows()
        if r == 0 or r == 1:
            selection.select_all()
        store, paths=selection.get_selected_rows()
        return [ store.get_value (store.get_iter(p), COLUMN_ELEMENT) for p in paths ]

    def set_elements(self, elements):
        model=self.build_model(elements)
        self.widget.treeview.set_model(model)
        self.model = model
        self.elements=elements

    def build_model(self, elements):
        """Build the ListStore containing the data.

        Columns: element, content (title), type, id
        """
        l=Gtk.ListStore(object, str, str, str)
        if not elements:
            return l
        for e in elements:
            l.append( (e,
                       self.controller.get_title(e),
                       helper.get_type(e),
                       getattr(e, 'id', str(e))) )
        return l

    def csv_export(self, name=None):
        if name is None:
            name=dialog.get_filename(title=_("Export data to file..."),
                                     default_file="advene_data.csv",
                                     action=Gtk.FileChooserAction.SAVE,
                                     button=Gtk.STOCK_SAVE)
        if name is None:
            return True
        try:
            f=open(name, 'w', encoding='utf-8')
        except IOError as e:
            dialog.message_dialog(label=_("Error while exporting data to %(filename)s: %(error)s"
                                          % {
                                              'filename': name,
                                              'error': str(e),
                                          }),
                                  icon=Gtk.MessageType.ERROR)
        w=csv.writer(f)
        tv=self.widget.treeview
        store, paths=tv.get_selection().get_selected_rows()
        source=[ store.get_iter(p) for p in paths ]
        if not source:
            source=tv.get_model()
        w.writerow( (_("Element title"), _("Element type"), _("Element id")) )
        for r in source:
            w.writerow( (str(r[COLUMN_CONTENT]).encode('utf-8'), str(r[COLUMN_TYPE]).encode('utf-8'), r[COLUMN_ID]) )
        f.close()
        self.log(_("Data exported to %s") % name)

    def build_widget(self):
        tree_view = Gtk.TreeView(self.model)

        select = tree_view.get_selection()
        select.set_mode(Gtk.SelectionMode.MULTIPLE)

        tree_view.connect('button-press-event', self.tree_view_button_cb)
        tree_view.connect('row-activated', self.row_activated_cb)
        # Deactivate starting search by simply typing. Users have to use the search shortcut (Control-F)
        tree_view.set_enable_search(False)
        #tree_view.set_search_column(COLUMN_CONTENT)

        def search_content(model, column, key, it):
            if key in model.get_value(it, COLUMN_CONTENT):
                return False
            return True

        tree_view.set_search_equal_func(search_content)

        columns={}
        for (name, label, col) in (
                ('title', _("Title"), COLUMN_CONTENT),
                ('type', _("Type"), COLUMN_TYPE),
                ('id', _("Id"), COLUMN_ID) ):
            columns[name]=Gtk.TreeViewColumn(label, Gtk.CellRendererText(), text=col)
            columns[name].set_reorderable(True)
            columns[name].set_sort_column_id(col)
            tree_view.append_column(columns[name])

        self.model.set_sort_column_id(COLUMN_CONTENT, Gtk.SortType.ASCENDING)

        # Resizable columns: title, type
        for name in ('title', 'type'):
            columns[name].set_sizing(Gtk.TreeViewColumnSizing.FIXED)
            columns[name].set_resizable(True)
            columns[name].set_min_width(40)
        columns['title'].set_expand(True)

        sw = Gtk.ScrolledWindow()
        sw.set_policy(Gtk.PolicyType.AUTOMATIC, Gtk.PolicyType.AUTOMATIC)
        sw.add(tree_view)

        sw.treeview = tree_view

        # The widget can receive drops
        def drag_received_cb(widget, context, x, y, selection, targetType, time):
            """Handle the drop of an annotation type.
            """
            if Gtk.drag_get_source_widget(context).is_ancestor(self.widget):
                # Ignore drops from our own widget
                return False

            if targetType == config.data.target_type['annotation']:
                sources=[ self.controller.package.annotations.get(uri) for uri in str(selection.get_data(), 'utf8').split('\n') ]
                if sources:
                    self.set_elements(sources)
                return True
            elif targetType == config.data.target_type['annotation-type']:
                sources=[ self.controller.package.annotationTypes.get(uri) for uri in str(selection.get_data(), 'utf8').split('\n') ]
                if sources:
                    self.set_elements(sources[0].annotations)
                return True
            return False

        tree_view.connect('drag-data-received', drag_received_cb)
        tree_view.drag_dest_set(Gtk.DestDefaults.MOTION |
                                Gtk.DestDefaults.HIGHLIGHT |
                                Gtk.DestDefaults.ALL,
                                config.data.get_target_types('annotation', 'annotation-type'),
                                Gdk.DragAction.COPY | Gdk.DragAction.LINK | Gdk.DragAction.MOVE)

        return sw

    def get_selected_nodes(self, with_path=False):
        """Return the currently selected nodes.
        """
        selection = self.widget.treeview.get_selection()
        store, paths = selection.get_selected_rows()
        if with_path:
            return [ (store.get_value(store.get_iter(p), COLUMN_ELEMENT), p) for p in paths ]
        else:
            return [ store.get_value(store.get_iter(p), COLUMN_ELEMENT) for p in paths ]

    def get_selected_node(self, with_path=False):
        """Return the currently selected node.

        None if no node is selected or multiple nodes are selected.
        """
        nodes = self.get_selected_nodes(with_path)
        if len(nodes) != 1:
            return None
        else:
            return nodes[0]

    def debug_cb (self, *p, **kw):
        logger.debug("Parameters: %s\nkw: %s", str(p), str(kw))

    def row_activated_cb(self, widget, path, view_column):
        """Edit the element on Return or double click
        """
        el = self.get_selected_node ()
        if el  is not None:
            self.controller.gui.edit_element(el)
            return True
        return False

    def tree_view_button_cb(self, widget=None, event=None):
        retval = False
        button = event.button
        x = int(event.x)
        y = int(event.y)

        if button == 3 or button == 2:
            if event.get_window() is widget.get_bin_window():
                model = self.model
                t = widget.get_path_at_pos(x, y)
                if t is not None:
                    path, col, cx, cy = t
                    it = model.get_iter(path)
                    node = model.get_value(it,
                                           COLUMN_ELEMENT)
                    widget.get_selection().select_path (path)
                    if button == 3:
                        menu = advene.gui.popup.Menu(node, controller=self.controller)
                        menu.popup()
                        retval = True
                    elif button == 2:
                        # Expand all children
                        widget.expand_row(path, True)
                        retval=True
        return retval


import json
import logging
import os
import sys
import time
from datetime import datetime
from logging import handlers
from random import randint
from websocket import create_connection

def checkaskfiles():
        global zeroaskoffset
        if os.path.isfile(ramdiskpath+'/'+market+'/asks/'+askprice) == True:
                fp = open(ramdiskpath+'/'+market+'/asks/'+askprice)
                line = fp.readline()
                fname = line.strip('\n\r').split(sep)
                fp.close()
                faskoffset = fname[0]
                fasksize = fname[1]
        else:
                faskoffset = 0
        if ( os.path.isfile(ramdiskpath+'/'+market+'/asks/'+askprice) == False and int(askoffset) > zeroaskoffset ) or int(askoffset) > int(faskoffset):
                if asksize == '0':
                        if os.path.isfile(ramdiskpath+'/'+market+'/asks/'+askprice) == True:
                                os.remove(ramdiskpath+'/'+market+'/asks/'+askprice)
                                zeroaskoffset = int(askoffset)
                else:
                        fp = open(ramdiskpath+'/'+market+'/asks/'+askprice, "w")
                        fp.write(askoffset+' '+asksize+' '+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+'\n')
                        fp.close()
                logger.info(datetime.now().strftime("%Y-%m-%d %H:%M:%S")+' Updated '+ramdiskpath+'/'+market+'/asks/'+askprice+': '+str('('+asksize+')').ljust(10)+' '+askoffset)

def checkbidfiles():
        global zerobidoffset
        if os.path.isfile(ramdiskpath+'/'+market+'/bids/'+bidprice) == True:
                fp = open(ramdiskpath+'/'+market+'/bids/'+bidprice)
                line = fp.readline()
                fname = line.strip('\n\r').split(sep)
                fp.close()
                fbidoffset = fname[0]
                fbidsize = fname[1]
        else:
                fbidoffset = 0
        if ( os.path.isfile(ramdiskpath+'/'+market+'/bids/'+bidprice) == False and int(bidoffset) > zerobidoffset ) or int(bidoffset) > int(fbidoffset):
                if bidsize == '0':
                        if os.path.isfile(ramdiskpath+'/'+market+'/bids/'+bidprice) == True:
                                os.remove(ramdiskpath+'/'+market+'/bids/'+bidprice)
                                zerobidoffset = int(bidoffset)
                else:
                        fp = open(ramdiskpath+'/'+market+'/bids/'+bidprice, "w")
                        fp.write(bidoffset+' '+bidsize+' '+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+'\n')
                        fp.close()
                logger.info(datetime.now().strftime("%Y-%m-%d %H:%M:%S")+' Updated '+ramdiskpath+'/'+market+'/bids/'+bidprice+': '+str('('+bidsize+')').ljust(10)+' '+bidoffset)

def openconnection():
        global ws
        global askprice
        global askoffset
        global asksize
        global bidprice
        global bidoffset
        global bidsize
#       ws = create_connection("wss://api.stage.dydx.exchange/v3/ws")
        ws = create_connection("wss://api.dydx.exchange/v3/ws")
        api_data = {
                "type": "subscribe",
                "channel": "v3_orderbook",
                "id": market,
                "includeOffsets": True
        }
        ws.send(json.dumps(api_data))
        api_data = ws.recv()
        api_data = json.loads(api_data)
        print(api_data)
        api_data = ws.recv()
        api_data = json.loads(api_data)
        asks = api_data['contents']['asks']
        bids = api_data['contents']['bids']
        for askitem in asks:
                askprice = askitem['price']
                askoffset = askitem['offset']
                asksize = askitem['size']
                checkaskfiles()
                checkwidth('price', len(askprice))
                checkwidth('size', len(asksize))
        for biditem in bids:
                bidprice = biditem['price']
                bidoffset = biditem['offset']
                bidsize = biditem['size']
                checkbidfiles()
                checkwidth('price', len(bidprice))
                checkwidth('size', len(bidsize))

def checkwidth(elementname, elementsize):
        global maxwidthprice
        global maxwidthsize
        if elementname == 'price' and elementsize > maxwidthprice:
                fp = open(ramdiskpath+'/'+market+'/maxwidth'+elementname, "w")
                fp.write(str(elementsize)+'\n')
                fp.close()
                maxwidthprice = elementsize
        elif elementname == 'size' and elementsize > maxwidthsize:
                fp = open(ramdiskpath+'/'+market+'/maxwidth'+elementname, "w")
                fp.write(str(elementsize)+'\n')
                fp.close()
                maxwidthsize = elementsize

print(datetime.now().strftime("%Y-%m-%d %H:%M:%S")+' dydxob.py')
logger = logging.getLogger("Rotating Log")
logger.setLevel(logging.INFO)
if sys.platform == "linux" or sys.platform == "linux2":
        # linux
        ramdiskpath = '/mnt/ramdisk'
elif sys.platform == "darwin":
        # OS X
        ramdiskpath = '/Volumes/RAMDisk'
sep = " "

if len(sys.argv) < 2:
        market = 'BTC-USD'
else:
        market = sys.argv[1]
handler = logging.handlers.RotatingFileHandler(ramdiskpath+'/dydxob'+market+'.log',
        maxBytes = 2097152,
        backupCount = 4
)
logger.addHandler(handler)

if os.path.isdir(ramdiskpath) == False:
        print('Error: Ramdisk', ramdiskpath, 'not mounted')
        sys.exit()
if os.path.ismount(ramdiskpath) == False:
        print('Warning:', ramdiskpath, 'is not a mount point')
if os.path.isdir(ramdiskpath+'/'+market+'/asks') == False:
        os.system('mkdir -p '+ramdiskpath+'/'+market+'/asks')
if os.path.isdir(ramdiskpath+'/'+market+'/bids') == False:
        os.system('mkdir -p '+ramdiskpath+'/'+market+'/bids')

maxwidthprice = 0
maxwidthsize = 0
zeroaskoffset = 0
zerobidoffset = 0
openconnection()
while True:
        try:
                api_data = ws.recv()
                api_data = json.loads(api_data)
                asks = api_data['contents']['asks']
                bids = api_data['contents']['bids']
                askoffset = api_data['contents']['offset']
                bidoffset = api_data['contents']['offset']
                if asks != []:
                        for askitem in asks:
                                askprice = askitem[0]
                                asksize = askitem[1]
                                checkaskfiles()
                                checkwidth('price', len(askprice))
                                checkwidth('size', len(asksize))
                if bids != []:
                        for biditem in bids:
                                bidprice = biditem[0]
                                bidsize = biditem[1]
                                checkbidfiles()
                                checkwidth('price', len(bidprice))
                                checkwidth('size', len(bidsize))
        except KeyboardInterrupt:
                ws.close()
                sys.exit(0)
        except Exception as error:
                print(datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "WebSocket message failed (%s)" % error)
                ws.close()
                time.sleep(1)
                try:
                        os.system('rm -rf '+ramdiskpath+'/'+market+'/asks/*')
                        os.system('rm -rf '+ramdiskpath+'/'+market+'/bids/*')
                        openconnection()
                except Exception as error:
                        print(datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "WebSocket message failed (%s)" % error)
                        ws.close()
                        time.sleep(randint(1,10))


import sys, os
sys.path.insert(1, os.path.join(sys.path[0], '..'))
import importlib
import matplotlib
importlib.reload(matplotlib)
matplotlib.use('Qt4Agg')
import vispy
import numpy as np
import random
from PyQt4.QtGui import *
from PyQt4 import QtCore, QtGui

from vispy import app, gloo
from vispy.util.transforms import perspective, translate, rotate
from vispy.geometry import meshdata as md
from vispy.geometry import generation as gen

from Deadweight.ShapeFileHandler import ShapeFileHandler
from Core.LasFileHandler import LasFileHandler

from Utils.pca_utils import *

import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt4agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QT

import Gui

from Core import DataHandler

class PlotObject():

    def __init__(self, parent=None):

        # a figure instance to plot on
        self.figure = plt.figure()

        # this is the Canvas Widget that displays the `figure`
        # it takes the `figure` instance as a parameter to __init__
        self.canvas = FigureCanvas(self.figure)

        # this is the Navigation widget
        # it takes the Canvas widget and a parent
        # self.toolbar = NavigationToolbar(self.canvas, self)
        self.toolbar = NavigationToolbar2QT(self.canvas, None)
        # Just some button connected to `plot` method
        self.button = QtGui.QPushButton('Plot')
        self.button.clicked.connect(self.plot)

        # set the layout
        self.layout = QtGui.QVBoxLayout()
        self.layout.addWidget(self.canvas)
        self.layout.addWidget(self.toolbar)
        self.layout.addWidget(self.button)
        #self.setLayout(layout)
        self.plot()

    def plot(self):
        ''' plot some random stuff '''
        # random data
        data = [random.random()*10 for i in range(10)]

        # create an axis
        ax = self.figure.add_subplot(111)

        # discards the old graph
        ax.hold(False)

        # plot data
        ax.plot(data, '*-')

        # refresh canvas
        self.canvas.draw()

class PipelineWidget(QWidget):

    def __init__(self, parent, data):

        QWidget.__init__(self)

        self.__parent__ = parent
        self.__data_handler = data
        self.current_data = {}
        self.setup_gui()
        self.resize(250, 150)

    def set_slice_from_property(self, property_id):
        self.current_data['bfps'], self.current_data['points'], self.current_data['solids'], self.current_data['property_area'], self.current_data['mbb'], self.current_data['top_dogs'] = self.__data_handler.get_slice_from_property(property_id)
        #self.__offset = np.mean(self.__property_area.points,0)

        self.current_data['top_dog_points'] = [self.current_data['points'][i] for i in range(len(self.current_data['bfps'])) if self.current_data['bfps'][i].id in self.current_data['top_dogs'] ]

    def setup_gui(self):
        #DATA GROUP
        self.data_group = QGroupBox("Data")
        self.data_layout = QVBoxLayout()
        self.data_group.setMaximumWidth(250)
        self.data_group.setLayout(self.data_layout)

        #self.bfp_checkbox = QCheckBox("Bfp blocks")
        #self.bfp_checkbox.setChecked(False)
        #self.bfp_checkbox.stateChanged.connect(lambda: self.bfps(self.bfp_checkbox.isChecked() ))

        #RENDER GROUP
        self.rendering_group = QGroupBox("Rendering")
        self.rendering_layout = QVBoxLayout()
        self.rendering_group.setMaximumWidth(250)
        self.rendering_group.setLayout(self.rendering_layout)

        # render polygons
        self.render_polygons_checkbox = QCheckBox("Render polygons?")
        self.render_polygons_checkbox.setChecked(True)
        self.render_polygons_checkbox.stateChanged.connect(lambda: self.setDrawPolygons(self.render_polygons_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.render_polygons_checkbox)
        # transparent checkbox
        self.transparent_checkbox = QCheckBox("Transparent polygons")
        self.transparent_checkbox.setChecked(False)
        self.transparent_checkbox.stateChanged.connect(lambda: self.transp(self.transparent_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.transparent_checkbox)
        # top dogs checkbox
        self.topdogs_checkbox = QCheckBox("Top dogs only")
        self.topdogs_checkbox.setChecked(False)
        self.topdogs_checkbox.stateChanged.connect(lambda: self.dogs(self.topdogs_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.topdogs_checkbox)

        # render points
        self.render_points_checkbox = QCheckBox("Render points?")
        self.render_points_checkbox.setChecked(True)
        self.render_points_checkbox.stateChanged.connect(lambda: self.setDrawPoints(self.render_points_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.render_points_checkbox)
        # reset points
        self.reset_points_button = QPushButton("reset points")
        self.reset_points_button.clicked.connect(lambda: self.reset_points())
        self.rendering_layout.addWidget(self.reset_points_button)
        #Colors
        self.colors_checkbox = QCheckBox("Debug colors")
        self.colors_checkbox.setChecked(False)
        self.colors_checkbox.stateChanged.connect(lambda: self.colors(self.colors_checkbox.isChecked()))
        self.rendering_layout.addWidget(self.colors_checkbox)

        #CLUSTERING GROUP
        self.cluster_group = QGroupBox("Clustering")
        self.cluster_layout = Gui.ClusterLayout(self)
        self.cluster_group.setMaximumWidth(300)
        self.cluster_group.setLayout(self.cluster_layout)

        #PLANE FITTING GROUP
        self.plane_group = QGroupBox("Plane Fitting")
        self.plane_layout = Gui.PlaneFittingLayout(self)
        self.plane_group.setMaximumWidth(200)
        self.plane_group.setLayout(self.plane_layout)

        #PLANE FITTING GROUP
        self.features_group = QGroupBox("Roof features creation")
        self.features_layout = Gui.FeatureExtractionLayout(self)
        self.features_group.setMaximumWidth(300)
        self.features_group.setLayout(self.features_layout)

        self.parent_layout = QHBoxLayout()
        self.pipeline_layout = QHBoxLayout()

        #self.pipeline_layout.addWidget(self.data_group)
        self.pipeline_layout.addWidget(self.rendering_group)
        self.pipeline_layout.addWidget(self.cluster_group)
        self.pipeline_layout.addWidget(self.plane_group)
        self.pipeline_layout.addWidget(self.features_group)

        self.parent_layout.setAlignment(QtCore.Qt.AlignLeft)
        self.parent_layout.addLayout(self.pipeline_layout)

        self.setLayout(self.parent_layout)

    def setDrawPoints(self, bool):
        self.__parent__.canvas_3d.render_params['draw_points'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def setDrawPolygons(self, bool):
        self.__parent__.canvas_3d.render_params['draw_polygons'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def setRegions(self, regions, topregions):
        self.current_data['regions'] = regions
        self.current_data['top_dog_regions'] = topregions
        self.__parent__.canvas_3d.regions()

    def reset_points(self):
        self.__parent__.canvas_3d.points()

    def setPlanes(self, planes, topplanes):
        self.current_data['planes'] = planes
        self.current_data['top_dog_planes'] = topplanes
        self.__parent__.canvas_3d.planes()

    def colors(self, value):
        self.__parent__.canvas_3d.render_params['debug_colors'] = value

    def setFeatures(self, features, topfeatures):
        self.current_data['features'] = features
        self.current_data['top_dog_features'] = topfeatures

        # import mpl_toolkits.mplot3d as a3
        # import matplotlib.colors as colors
        # import pylab as pl
        # import scipy as sp
        # import triangle
        # ax = a3.Axes3D(pl.figure())
        # colors = [[1,0,0],[0,1,0], [0,0,1], [0.65, 0.65, 0], [0, 0.65, 0.65], [0.65, 0, 0.65], [0.43, 0.43, 0.43]]
        #
        #
        # f = features
        # if f == None:
        #     f = topfeatures
        #
        # cid = 0
        # for feature_group in f:
        #     for feature in feature_group:
        #         roofPolys = feature["roof"].get_triangles_as_polygons()
        #         wallPolys = feature["walls"].get_triangles_as_polygons()
        #         #vertexdata = feature.get_gl_vertex_data()
        #
        #         tri = a3.art3d.Poly3DCollection([poly.exterior.coords for poly in wallPolys])
        #         tri.set_facecolor([0.8,0.8, 0.8])
        #         ax.add_collection3d(tri)
        #         tri = a3.art3d.Poly3DCollection([poly.exterior.coords for poly in roofPolys])
        #         tri.set_facecolor(colors[cid%7])
        #
        #         ax.add_collection3d(tri)
        #     cid += 1
        # ax.view_init(elev=39, azim=-55)
        # ax.set_xlim3d(-30, 30)
        # ax.set_ylim3d(-30, 30)
        # ax.set_zlim3d(0, 60)
        # pl.show()

        self.__parent__.canvas_3d.features()

    def transp(self, bool):
        self.__parent__.canvas_3d.render_params['transparent'] = bool

    def dogs(self, bool):
        self.__parent__.canvas_3d.render_params['top_dogs_only'] = bool
        self.__parent__.canvas_3d.updateCanvas()

    def bfps(self, bool):
        self.__parent__.canvas_3d.draw_model_transparent = bool

    def render_points(self, attr, value):
        #self.__parent__.canvas_3d.setup_bfps()
        pass

class MainWindow(QMainWindow):
    """
        this is the mainWondow for the program gui
    """

    def __init__(self, data_handler):

        QMainWindow.__init__(self)

        self.__data_handler = data_handler

        self.resize(1400, 840)
        self.setWindowTitle('gui main window')

        self.pipelineWidget = PipelineWidget(self, data_handler)

        self.canvas_map = Gui.Map2D(data_handler, self)
        self.canvas_map.create_native()
        self.canvas_map.native.setParent(self)

        self.canvas_3d = Gui.Canvas3D(data_handler, self, self.pipelineWidget.current_data)
        self.canvas_3d.create_native()
        self.canvas_3d.native.setParent(self)

        self.listWidget = Gui.BuildingListWidget(self, data_handler)

        self.setupLayout()
        #self.__focus_building_id = 0

        self.set_slice(1592) #1593
        self.show()

    #Is called from the property list
    def set_slice(self, id):
        print("selected property id: ", id)
        self.pipelineWidget.set_slice_from_property(id)
        self.canvas_3d.setCanvasData()

    @property
    def focus_building_id(self):
        return self.__focus_building_id

    @focus_building_id.setter
    def focus_building_id(self, id_to_focus):
        self.__focus_building_id = id_to_focus
        self.canvas_map.set_selected(self.__focus_building_id)
        #self.canvas_3d.set_selected(self.__focus_building_id)

    def setupLayout(self):

        self.verticalSplitter = QSplitter(QtCore.Qt.Vertical)

        self.horizontalSplitter = QSplitter(QtCore.Qt.Horizontal)

        self.mapListSplitter = QSplitter(QtCore.Qt.Vertical)

        self.mapListSplitter.addWidget(self.listWidget)
        self.mapListSplitter.addWidget(self.canvas_map.native)

        self.horizontalSplitter.addWidget(self.mapListSplitter)

        self.horizontalSplitter.addWidget(self.canvas_3d.native)

        self.verticalSplitter.addWidget(self.horizontalSplitter)
        self.verticalSplitter.addWidget(self.pipelineWidget)
        self.setCentralWidget(self.verticalSplitter)

if __name__ == "__main__":
    appQt = QApplication(sys.argv)

    #value = Vestibule(None)

    #filename = ("by_get_shp/by_get")

    #shapes = ShapeFileHandler(filename)

    #las_dir_path = os.path.dirname(os.path.realpath(__file__))
    #las_dir_path = las_dir_path + '/datafromndr'
    #las_file_path = las_dir_path + "/norrkoping.las"
    #las = LasFileHandler(las_file_path,8)

    #las.read_points()
    #las.read_points()

    #mw = MainWindow(shapes, las) #gui main

    buildings_path = "by_get_shp/by_get"
    propertys_path = "ay_get_shp/ay_get"
    las_path = "datafromndr/norrkoping.las"
    data_handler = DataHandler(buildings_path,propertys_path,las_path)

    mw = MainWindow(data_handler)

    sys.exit(appQt.exec_())


from ..common.boundingbox import BoundingBox
import cv2
import numpy as np

""" 各自圖片進行撿均值（突顯該 圖像 重點） """
def subtract_avg(img):

    trg_img = img
    for i in range(3):
        avg=np.average(trg_img[:,:,i])
        trg_img[:,:,i]-=avg
    return trg_img

""" 整個數據集進行撿均值（突顯該 類別 重點） """
def subtract_offset(img, offset=( 103.939, 116.779, 123.68 )):
    trg_img = img
    for i in range(3):
        trg_img[:,:,i]=img[:,:,i]-offset[i]
    return trg_img

""" caffe mode of image processing """
def preprocess(img, input_shape, letter_box=False):
    """Preprocess an image before TRT YOLO inferencing.
    # Args
        img: int8 numpy array of shape (img_h, img_w, 3)
        input_shape: a tuple of (H, W)
        letter_box: boolean, specifies whether to keep aspect ratio and
                    create a "letterboxed" image for inference
    # Returns
        preprocessed img: float32 numpy array of shape (3, H, W)
    """
    if letter_box:
        img_h, img_w, _ = img.shape
        new_h, new_w = input_shape[0], input_shape[1]
        offset_h, offset_w = 0, 0
        if (new_w / img_w) <= (new_h / img_h):
            new_h = int(img_h * new_w / img_w)
            offset_h = (input_shape[0] - new_h) // 2
        else:
            new_w = int(img_w * new_h / img_h)
            offset_w = (input_shape[1] - new_w) // 2
        resized = cv2.resize(img, (new_w, new_h))
        img = np.full((input_shape[0], input_shape[1], 3), 127, dtype=np.uint8)
        img[offset_h:(offset_h + new_h), offset_w:(offset_w + new_w), :] = resized
    else:
        img = cv2.resize(img, (input_shape[1], input_shape[0]))

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img.transpose((2, 0, 1)).astype(np.float32)
    img /= 255.0
    return img

def postprocess(output, img_w, img_h, input_shape, conf_th=0.8, nms_threshold=0.5, letter_box=False):
    """
    (-1,200), (-1,200,4), (-1,200), (-1,200)
    num_detections: A [batch_size] tensor containing the INT32 scalar indicating the number of valid detections per batch item. It can be less than keepTopK. Only the top num_detections[i] entries in nmsed_boxes[i], nmsed_scores[i] and nmsed_classes[i] are valid
    nmsed_boxes: A [batch_size, keepTopK, 4] float32 tensor containing the coordinates of non-max suppressed boxes
    nmsed_scores: A [batch_size, keepTopK] float32 tensor containing the scores for the boxes
    nmsed_classes: A [batch_size, keepTopK] float32 tensor containing the classes for the boxes
    """
    (detections, bboxes, scores, labels) = output

    results= []
    for idx, det in enumerate(detections):
        # print(scores[idx], '\t')
        x1, y1, x2, y2 = map(float, bboxes[idx]) # x1, y1, x2, y2
        x1, y1, x2, y2 = x1*img_w, y1*img_h, x2*img_w, y2*img_h
        if scores[idx]>=conf_th:
            results.append(BoundingBox(labels[idx], scores[idx], x1, x2, y1, y2, img_h, img_w))
    return results


import statsmodels.api as sm
import matplotlib.pyplot as plt
import numpy as np
import itertools
from scipy.stats import skewnorm

means = [100, 105]
std_devs = [5, 6]

num_to_sim = 500

data = [
    np.random.randn(num_to_sim) * s + mu for mu, s in zip(means, std_devs)
] + [skewnorm.rvs(10, size=num_to_sim)*std_devs[0] + means[0]]

plt.hist(data[0])
plt.xlabel('Value')
plt.ylabel('Count')
plt.savefig('imgs/basic_histogram.png')

plt.figure()
plt.hist(data[0])
plt.hist(data[1], alpha=0.5)
# plt.hist(data[2], alpha=0.5)

plt.xlabel('Value')
plt.ylabel('Count')
plt.savefig('imgs/two_distributions.png')

pp = sm.ProbPlot(data[0])
pp2 = sm.ProbPlot(data[1])
pp3 = sm.ProbPlot(data[2])
fig = pp.probplot()
pp2.probplot(ax=fig.gca())
pp3.probplot(ax=fig.gca())
plt.grid(True)

colors = itertools.cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])

for l in fig.gca().get_lines():
    c = next(colors)
    l.set_markerfacecolor(c)
    l.set_color(c)

plt.tight_layout()
plt.savefig('imgs/probplot.png')


from sklearn.cluster import  MiniBatchKMeans, KMeans, DBSCAN
import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import pandas as pd
import math
from pylab import *
from scipy.spatial.distance import cdist
from scipy import spatial
from mpl_toolkits.mplot3d import Axes3D
import time

from utils.logistics_functions import *
from utils.Parameters_class import Parameters_values

Parameters = Parameters_values()

class logistics_model():
	def __init__(self, path, n_regen, n_collection, logistics, scenario):
		self.path =path
		self.n_regen = n_regen
		self.n_collection = n_collection
		self.logistics = logistics
		self.scenario = scenario

	def return_dataframe_buildings(self):
		building_virtual_buildings_df = pd.read_csv(self.path)
		building_SF_points=building_virtual_buildings_df[['lat_lat','lon_lon']]
		return building_SF_points

	def clustering_regen(self):
		k_means_labels_regen, k_means_cluster_centers_regen = clustering(self.return_dataframe_buildings(), self.n_regen)
		return k_means_labels_regen, k_means_cluster_centers_regen

	def clustering_regen_grid(self):
		building_SF_points = self.return_dataframe_buildings()
		nx = int(sqrt(self.n_regen))
		ny = int(ceil(self.n_regen/nx))
		cluster_centers_regen = make_grid_points(self.n_regen, nx, ny)
		cluster_centers_regen_df=pd.DataFrame(cluster_centers_regen, columns=['lat', 'lon'])
		tree = spatial.KDTree(list(zip(cluster_centers_regen_df['lat'], cluster_centers_regen_df['lon'])))
		building_SF_points_array = list(zip(building_SF_points['lat_lat'],building_SF_points['lon_lon']))
		k_means_labels_regen = tree.query(building_SF_points_array)[1]

		return k_means_labels_regen, cluster_centers_regen

	def clustering_regen_random(self):
		building_SF_points = self.return_dataframe_buildings()
		cluster_centers_regen = make_random_points(self.n_regen)
		cluster_centers_regen_df=pd.DataFrame(cluster_centers_regen, columns=['lat', 'lon'])
		tree = spatial.KDTree(list(zip(cluster_centers_regen_df['lat'], cluster_centers_regen_df['lon'])))
		building_SF_points_array = list(zip(building_SF_points['lat_lat'],building_SF_points['lon_lon']))
		k_means_labels_regen = tree.query(building_SF_points_array)[1]

		return k_means_labels_regen, cluster_centers_regen

	def clustering_collection(self, k_means_cluster_centers_regen):
		k_means_labels_collection, k_means_cluster_centers_collection = clustering(k_means_cluster_centers_regen, self.n_collection)
		return k_means_labels_collection, k_means_cluster_centers_collection

	def logistics_distances(self):
		building_virtual_buildings_df = pd.read_csv(self.path)

		if self.logistics == 'optimal':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen()

		elif self.logistics == 'grid':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen_grid()

		elif self.logistics == 'random':
			k_means_labels_regen, k_means_cluster_centers_regen = self.clustering_regen_random()

		else:
			print ("Error type of logistics was not specified")

		time_now = time.time()

		print ("Start calculating distances...")
		distance_regeneration = find_distance_regeneration_scheduled(building_virtual_buildings_df, k_means_labels_regen, self.scenario)
		distance_regeneration['total_dist_m_y'] = distance_regeneration['total_dist_m']*365/Parameters.time_between_catridge_regeneration_2
		time_end = time.time() - time_now
		print ("calc distances took time %s" %time_end)

		distance_collection = find_distance_collection(k_means_cluster_centers_regen)*Parameters.collection_times_per_year
		return distance_regeneration, distance_collection


import torch
from interactive_navigation.utils.train_mask_rcnn import get_model_instance_segmentation
from PIL import Image
import numpy as np
import cv2

if __name__ == "__main__":
    root_dir = "storage/heuristic_corner_detector_test_scene"
    OBJ = sorted(["ArmChair", "DogBed", "Box", "Chair", "Desk", "DiningTable", "SideTable", "Sofa",
                  "Stool", "Television", "Pillow", "Bread", "Apple", "AlarmClock", "Lettuce",
                  "GarbageCan", "Laptop", "Microwave", "Pot", "Tomato"])
    SIZE = (400, 400)
    COLOR = [(0, 0, 0),
             (255, 0, 0),
             (255, 128, 0),
             (255, 255, 0),
             (128, 255, 0),
             (0, 255, 0),
             (0, 255, 128),
             (0, 255, 255),
             (0, 128, 255),
             (0, 0, 255),
             (127, 0, 255),
             (255, 0, 255),
             (255, 0, 127),
             (128, 128, 128),
             (102, 102, 0),
             (102, 0, 0),
             (0, 102, 0),
             (0, 102, 102),
             (0, 0, 102),
             (102, 0, 102),
             (255, 255, 255)]

    gpu_id = 0
    loc = 'cuda:{}'.format(gpu_id)
    checkpoint = torch.load("storage/maskRcnn/model_9.pth", map_location=loc)
    model = get_model_instance_segmentation(21).cuda(gpu_id)
    model.load_state_dict(checkpoint["model"])
    model.eval()

    imgs = []
    for obj in OBJ:
        img = Image.open("{}/{}/obs.png".format(root_dir, obj))
        img = img.resize(SIZE)
        img = np.array(img) / 255.
        img = np.transpose(img, (2, 0, 1))
        img = torch.Tensor(img).to(gpu_id)
        imgs.append(img)
    output = model(imgs)
    for j, out in enumerate(output):
        labels = out["labels"].detach().cpu().numpy()
        scores = out["scores"].detach().cpu().numpy()
        masks = out["masks"].squeeze().detach().cpu().numpy()
        result = np.zeros((SIZE[0], SIZE[1], 3))
        for i, label in enumerate(labels):
            if label == (j + 1):
                if scores[i] > 0.5:
                    color = COLOR[label]
                    mask = masks[i]
                    idx = np.where(mask > 0.1)
                    result[idx] = color
        result = result.astype(np.uint8)
        img = Image.fromarray(result)
        img.save("{}/{}/maskrcnn_segmentation.png".format(root_dir, OBJ[j]))

        org_img = np.transpose(imgs[j].detach().cpu().numpy() * 255, (1, 2, 0)).astype(np.uint8)
        fin = cv2.addWeighted(result, 0.5, org_img, 0.85, 0)
        fin = Image.fromarray(fin)
        fin.save("{}/{}/maskrcnn_segmentation_blend.png".format(root_dir, OBJ[j]))

        org_img = Image.fromarray(org_img)
        org_img.save("{}/{}/maskrcnn_org_image.png".format(root_dir, OBJ[j]))

    xxx = 0


import numpy as np
import matplotlib.pyplot as plt
import torch

class control_env():
    def __init__(self, steps):
        self.num_agents = 3
        self.Qii = torch.tensor([[1.0, 0.0],[0.0, 1.0]])
        self.Rij = torch.tensor([[1.0, 0.0, 0.0],[0.0, 1.0 ,0.0],[0.0, 1.0, 1.0]])
        self.A = torch.tensor([[0.495, 0.09983], [0.09983, 0.495]])
        self.B = torch.tensor([[0, 0.1],[0, 0.9],[0, 0.8]])
        self.aij = torch.tensor([[0.0, 0.0, 0.0],[0.0, 0.0 ,0.0],[0.0, 1.0, 0.0]])
        self.bi = torch.tensor([[1],[1],[0]])
        self.n_steps = steps

    def forward(self, i, u = None, X_old = None):
        # Initialize the control environment with the given number of steps
        num_agents = self.num_agents
        Qii = self.Qii
        Rij = self.Rij
        A = self.A
        B = self.B
        aij = self.aij
        bi = self.bi
        n_steps = self.n_steps

        pi = torch.tensor(np.pi)

        def Ji(k, ei, ein, U, alpha, Qii, Rij):
            # Calculate the cost given the current inputs, state, and control errors
            def reward(e):
                reward = torch.dot(torch.matmul(e, Qii), e) + U[k] * Rij[k][k] * U[k]
                for i in range(num_agents):
                    reward = reward + U[i] * Rij[k][i] * U[i]
                return reward

            cost = reward(ei)
            # print(cost)
            return cost

        t = torch.linspace(0, 15, n_steps)
        X_ref = torch.zeros((2, 2, n_steps))
        X = torch.zeros((num_agents, 2, 2, n_steps))

        if X_old == None:
            X_old = torch.zeros((3,2))

        if u is None:
            u = ((torch.randn(num_agents)) * 0) - 2.5
            # print(u)

        X[:,0,:,i] = X_old

        track_error = torch.zeros((num_agents, 2))
        index_per = torch.zeros((num_agents))

        x0 = torch.tensor([[3*torch.sin(1/2*pi*t[i]), 2*torch.sin(1/2*pi*t[i] + pi/2)]]).T
        # Compute state at each time step
        X_ref[0,:,i] = x0.flatten()
        X_ref[1,:,i] = torch.matmul(A, X_ref[0,:,i])
        ei = torch.zeros((num_agents, 2))
        ein = torch.zeros((num_agents, 2))

        for j in range(num_agents):
             # Compute the new state vector based on the current input and state vector
            # X[j,0,:,i] = x0.flatten()
            X[j,1,:,i] = torch.matmul(A, X[j,0,:,i]) + B[j] * u[j]

        for k in range(num_agents):

            for j in range(num_agents):
                ei[k,:] = ei[k,:] + aij[k][j] * (X[k,0,:,i] - X[j,0,:,i])
                ein[k,:] = ein[k,:] + aij[k][j] * (X[k,1,:,i] - X[j,1,:,i])
            ei[k,:] = ei[k,:] + bi[k] * (X[k,0,:,i] - X_ref[0,:,i])
            ein[k,:] = ein[k,:] + bi[k] * (X[k,1,:,i] - X_ref[1,:,i])
            index_per[k] = Ji(k, ei[k,:], ein[k,:], u, 0.5, Qii, Rij)

            # print(track_error)

        return ei, index_per, u, X_ref[1,:,i], X[:,1,:,i]

if __name__ == "__main__":
    env = control_env(300)
    ei, index_per, u, X_ref, X = env.forward(1)
    fig = plt.figure()
    ax = fig.add_subplot(projection='3d')

    # Initialize lists to store data for visualization
    ref_1 = []
    ref_2 = []

    agents_1 = []
    agents_2 = []

    X_old = torch.zeros((3,2))

    # Loop over each time step
    for i in range(300):
        # Get outputs from environment
        u = torch.randn(3)
        ei, index_per, u, X_ref, X = env.forward(i, u, X_old)

        # Store relevant data for visualization
        ref_1.append(X_ref[0])
        ref_2.append(X_ref[1])
        for j in range(3):
            agents_1[j].append(X[j,0])
            agents_2[j].append(X[j,1])

        # Update X_old for next iteration
        X_old = X

    # Visualize data in the plot
    xs = ref_1
    ys = ref_2
    zs = list(range(0,300))
    ax.scatter(xs, ys, zs, marker='o')

    for j in range(3):
        xs = agents_1[j]
        ys = agents_2[j]
        zs = list(range(0,300))
        ax.scatter(xs, ys, zs, marker='o')

    # Add labels to the plot
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()

    # print(ei)
    # print(ei[1])
    itorchuts = torch.concatenate((ei.flatten(), u))
    # print(itorchuts)
    itorchuts = torch.tensor(itorchuts, dtype=torch.float).to("cpu")


import pygame

class Ship:

    def __init__(self, ai_settings, screen):
        self.screen = screen
        self.settings = ai_settings    # converting to the attribute for use in functions of class

        self.image = pygame.image.load('images/starship.png')
        self.rect = self.image.get_rect()
        self.screen_rect = screen.get_rect()

        self.rect.centerx = self.screen_rect.centerx
        self.rect.bottom = self.screen_rect.bottom

        self.center = float(self.rect.centerx)      # better change control for ship speed parameter

        self.moving_right = False
        self.moving_left = False

    def update(self):
        if self.moving_right and self.rect.right < self.screen_rect.right:
            self.rect.centerx += self.settings.ship_speed_factor
        if self.moving_left and self.rect.left > 0:     # this implementation allows you to set equal priority
            self.rect.centerx -= self.settings.ship_speed_factor

    def blitme(self):
        """draw the ship at its current location"""
        self.screen.blit(self.image, self.rect)


import sys
import argparse
import logging
from binaryninja import *

import analyzers.config
from analyzers.mliltracer import *

parser = argparse.ArgumentParser()
parser.add_argument("filename", help = "path to binary ninja bndb file")
parser.add_argument("--loglevel", help = "choose logging level for debugging", default = "INFO")
parser.add_argument("--logname", help = "choose log file name", default = "analysis.log")
parser.add_argument("--function_hooks", help = "provide path to json config")
args = parser.parse_args()

bv = binaryview.BinaryViewType.get_view_of_file(args.filename)

numeric_level = getattr(logging, args.loglevel.upper(), None)
logger = logging.getLogger()
logger.setLevel(numeric_level)

filehandler = logging.FileHandler(args.logname)
filehandler.setLevel(numeric_level)
logger.addHandler(filehandler)

consolehandler = logging.StreamHandler(sys.stdout)
consolehandler.setLevel(numeric_level)
logger.addHandler(consolehandler)

if args.function_hooks is not None:
    with open(args.function_hooks) as func_hooks:
        config.function_hooks = json.load(func_hooks)

# configure start of MySQL Signal Data
config.struct_offset = 0x28

# configure integer ranges to filter FP
config.ALLOW_VAL = 0xFFFE
config.LSR_RANGE = 0x20
config.MUL_RANGE = 0

arg_for_index = lambda x: "arg" + str(x + 1)

count = 0
logging.info("Starting analysis...")

# For all functions in binary, check if the parameter is a pointer to Signal

for func in bv.functions:
    for index, param in enumerate(func.parameter_vars):
        if param.type is not None and param.type.type_class == TypeClass.PointerTypeClass:
            if param.type.tokens[0].text == "Signal":

                args = dict()
                typeinfo, demangled_name = demangle_gnu3(bv.arch, func.name)
                function_name = get_qualified_name(demangled_name)

                # Set base address as 0 for Signal * argument for propagation
                args[arg_for_index(index)] = 0

                # In case a vulnerability is found during tracing, this flag is set to TRUE
                config.vuln = False

                func_trace = MLILTracer(bv, func.start)
                func_trace.set_function_args(args)
                func_trace.trace()
                del func_trace

                if config.vuln == True:
                    logging.info("[%d] %s" % (count, function_name))
                    count = count + 1


#!/usr/bin/env python3

import sys

sys.path.append('/root/gcamp_ros2_ws/src/gcamp_ros2_basic/posenet_pkg/posenet_pkg')

from geometry_msgs.msg import Pose
from sensor_msgs.msg import Image as Images
import rclpy
from rclpy.node import Node

import cv2
from cv_bridge import CvBridge
import numpy as np
from PIL import Image
from torch.backends import cudnn

from node_data_loader import get_loader
from node_solver import Solver

class PosenetNode(Node):

  def __init__(self):
    super().__init__('posenet_node')

    self.publisher = self.create_publisher(Pose, 'posenet_pytorch', 10)

    self.subscriber = self.create_subscription(Images, 'image', self.sub_callback, 10)

    self.br = CvBridge()

    self.get_logger().info('Node Started, Waiting.... \n')

  def sub_callback(self, msg):

    opencv_img = self.br.imgmsg_to_cv2(msg)
    cvt_img = cv2.cvtColor(opencv_img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(cvt_img)
    cudnn.benchmark = True
    data_loader = get_loader(image=pil_img,model='Resnet',mode='test', batch_size=1)

    solver = Solver(data_loader)

    pos, ori = solver.test()
    pose_msg = Pose()
    pose_msg.position.x = float(pos[0])
    pose_msg.position.y = float(pos[1])
    pose_msg.position.z = float(pos[2])
    pose_msg.orientation.x = float(ori[1])
    pose_msg.orientation.y = float(ori[2])
    pose_msg.orientation.z = float(ori[3])
    pose_msg.orientation.w = float(ori[0])

    self.get_logger().info(f'position [x:{pose_msg.position.x}, y:{pose_msg.position.y}, z:{pose_msg.position.z}]')
    self.get_logger().info(f'orientation [x:{pose_msg.orientation.x}, y:{pose_msg.orientation.y}, z:{pose_msg.orientation.z}, w:{pose_msg.orientation.w}]')
    self.publisher.publish(pose_msg)
    self.get_logger().info('====Done====')

def main(args=None):

  rclpy.init(args=args)
  posenet_node = PosenetNode()

  rclpy.spin(posenet_node)

  rclpy.shutdown()
  print("done")

if __name__ == '__main__':
  main()


# -*- coding: utf-8 -*-
"""
Created on Thu Feb  1 16:05:49 2018

@author: peter
"""
import logging
import time
import unittest

from CESAPI.command import *
from CESAPI.connection import Connection
import CESAPI.packet
from CESAPI.refract import *
from CESAPI.test import LTSimulator

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Base class for most of the other test cases. Used to setup and teardown
# the LTSimulator instance that the test cases use.
class RefractTestCase(unittest.TestCase):
    def __init(self):
        self.sim = None

    def setUp(self):
        self.sim = LTSimulator()
        self.sim.start()

    def tearDown(self):
        self.sim.stop()
        self.sim.join()

# Send an InitializeCT packet and verify that an InitializeRT packet was returned.
class SetRefractionIndexTestCase(RefractTestCase):
    def runTest(self):
        success = False
        connection = Connection()
        initrt = None
        try:
            connection.connect(host='localhost')
            command = CommandSync(connection)

            algorithm = AlgorithmFactory().refractionIndexAlgorithm(RI_ALG_Leica)
            SetRefractionIndex(command, algorithm)
            success = True
        except Exception as e:
            logger.info(e)
        finally:
            if connection != None:
                connection.disconnect()

        self.assertTrue(success)

if __name__ == '__main__':
    unittest.main()


# Download Video YouTube
#===============================================================
# GỌI HÀM menu_update() ĐỂ ĐÓNG SERVER TOOLS

# notification()

#===============================================================

import os
try:
	from pytube import YouTube
except:
	os.system('pip install pytube')
	os.system("cls")

try:
	from pytube import YouTube
	print("\n*******************************************************************************\n")
	urls = input('>>> Nhập những Url Video YouTube muốn tải: ')
	urls = " ".join(urls.split())
	urls = urls.split()

	def down(url):
		yt = YouTube(url)
		video = yt.streams.filter(progressive = True, file_extension = 'mp4').last().download()
		print(f'>>> Video đang ở {video}')

	for url in urls:
		try:
			down(url)
		except:
			print('>>> Download Fail!')

except KeyboardInterrupt:
    r = requests.get("https://raw.githubusercontent.com/HaLTools/Server_Test/master/theme.py")
    exec(r.text)


# coding=utf-8
from django.conf.urls import url
from . import views
from django.contrib.auth import views as auth_views

app_name = 'accounts'

urlpatterns = [

	# url(r'^gestemail/$', views.CreateGestEmail.as_view(), name='CreateGestEmail'),
	url(r'^profile/(?P<pk>\d+)/$', views.ProfileView.as_view(), name='profile'),
	url(r'^createuser/$', views.CreateUserView.as_view(), name="createuser",),
	url(r'^confirm/$', views.ConfirmView.as_view(), name="confirm",),
	# url(r'^login/$', auth_views.LoginView.as_view(template_name='accounts/login.html'), name="login",),
	url(r'^login/$', views.MyLoginView.as_view(), name="login",),
	url(r'^logout/$', auth_views.LogoutView.as_view(), name="logout",),
	url(r'^logout-then-login/$', auth_views.logout_then_login, name='logout_then_login'),
]


import unittest
import time
from selenium import webdriver

def webdriverChrome():
    options = webdriver.ChromeOptions()
    options.add_experimental_option('excludeSwitches', ['enable-logging'])
    return webdriver.Chrome(options=options)

class TestRegistration(unittest.TestCase):

    def setUp(self):
        self.browser = webdriverChrome()

    def tearDown(self):
        self.browser.quit()

    def registration(self, link):
        browser = self.browser
        browser.get(link)

        fields = {
            ".first_block .first": "Ivan",
            ".first_block .second": "Petrov",
            ".first_block .third": "ivan@gmail.com",
        }
        for selector, value in fields.items():
            browser.find_element_by_css_selector(selector).send_keys(value)

        browser.find_element_by_css_selector("button.btn").click()

        time.sleep(1)
        welcome_text = browser.find_element_by_tag_name("h1").text

        self.assertTrue("successfully" in welcome_text, "registration is failed")

    def test_registration1(self):
        self.registration("http://suninjuly.github.io/registration1.html")

    def test_registration2(self):
        self.registration("http://suninjuly.github.io/registration2.html")

if __name__ == "__main__":
    unittest.main()


import torch
import torch.optim as optim

import numpy as np
import time

import os
from datetime import datetime
from torch.utils.tensorboard import SummaryWriter

from Seq2Seq_model import Seq2Seq_VAE

from helper_utils import colorize
from training_functions import run_VAE

from helper_utils import random_trajectories

from LSL_utils import initial_region_points
from LSL_utils import find_good_region
from LSL_utils import set_param_grad_status

from LSL_utils import trajectory_diagnostics

import copy

import roboschool
import gym

def main(main_anneal='logistic',
         main_num_compression_iterations=30,
         main_num_logistic_compression_iterations=15,
         main_initial_num_epochs=3000,
         main_num_const_epochs=10,
         bootstrap="no", given_time="",
         final_model_file=None, old_final_model_file=None, argmax_only=False, small_steps=False,
         main_z_size=2, main_enc_hidden_size=12, main_dec_hidden_size=50, main_num_layers=1,
         main_bidir=True, main_use_layer_norm=True, main_use_ff_dropout=False, main_use_recurrent_dropout=False,
         main_ff_drop_prob=0.0, main_recurrent_drop_prob=0.0, main_dec_num_layers=3,
         main_Adam_learning_rate=5e-3, main_clip_gn=500.0, main_clip_gv=100.0, main_kl_tolerance=0.0,
         main_target_beta=1.0, main_logistic_scale=12.0, main_floor=0.5, main_num_grid_points=400,
         main_num_trajectories=240, main_batch_size=8, main_seq_len=20,
         main_env_name = "MyCartPoleNew-v0",
         main_torch_manual_seed=684, main_np_random_seed=684, main_env_seed=6328):

    torch.set_printoptions(precision=3, sci_mode=False)
    np.set_printoptions(precision=3, suppress=True)

    torch_manual_seed = main_torch_manual_seed
    np_random_seed = main_np_random_seed
    env_seed = main_env_seed

    torch.manual_seed(torch_manual_seed)
    np.random.seed(np_random_seed)

    env_name = main_env_name

    z_size = main_z_size
    enc_hidden_size = main_enc_hidden_size
    dec_hidden_size = main_dec_hidden_size
    num_layers = main_num_layers
    bidir = main_bidir
    use_layer_norm = main_use_layer_norm

    use_ff_dropout = main_use_ff_dropout
    use_recurrent_dropout = main_use_recurrent_dropout
    ff_drop_prob = main_ff_drop_prob
    recurrent_drop_prob = main_recurrent_drop_prob

    dec_num_layers = main_dec_num_layers

    seq_len = main_seq_len

    env = gym.make(env_name)
    env.seed(env_seed)
    obs_dim = env.observation_space.shape

    env_state_size = obs_dim[0]
    action_size = env.action_space.n

    Adam_learning_rate = main_Adam_learning_rate
    clip_gn = main_clip_gn
    clip_gv = main_clip_gv
    kl_tolerance = main_kl_tolerance
    target_beta = main_target_beta
    logistic_scale = main_logistic_scale

    anneal = main_anneal
    floor = main_floor

    num_grid_points = main_num_grid_points
    num_trajectories = main_num_trajectories
    batch_size = main_batch_size

    current_time = datetime.now().strftime('%b%d_%H-%M-%S')

    model_file = ('saved_models/model' + anneal + str(torch_manual_seed) + current_time + '.pt')

    num_compression_iterations = main_num_compression_iterations

    trajectory_actions, trajectory_actions_one_hot, trajectory_env_states = \
        random_trajectories(num_trajectories, seq_len, action_size, env_state_size, env=env)

    primitive_old_model_file = old_final_model_file
    primitive_model_file = final_model_file

    if bootstrap != "no":
        load_model = True
    else:
        load_model = False

    if load_model:

        print("loading model: " + primitive_model_file)
        model: Seq2Seq_VAE = torch.load(primitive_model_file)
        print(model)

        new_model = copy.deepcopy(model)

        if primitive_old_model_file is not None:
            print("loading model: " + primitive_old_model_file)
            model: Seq2Seq_VAE = torch.load(primitive_old_model_file)
            print(model)

        new_model.h2_bottom_seq_len = 7
        new_model.h2_top_seq_len = 10

        new_model.h3_bottom_seq_len = 7
        new_model.h3_top_seq_len = 10

    else:

        print("creating new model")

        num_hierarchies = 3

        model = Seq2Seq_VAE(action_size, env_state_size, z_size, enc_hidden_size, dec_hidden_size,
                            num_layers, dec_num_layers=dec_num_layers, seq_len=seq_len,
                            bidirectional_encoder=bidir,
                            forget_bias=1.0,
                            use_ff_dropout=use_ff_dropout, ff_drop_prob=ff_drop_prob, same_ff_mask=True,
                            use_recurrent_dropout=use_recurrent_dropout, recurrent_drop_prob=recurrent_drop_prob,
                            same_recurrent_mask=True,
                            use_layer_norm=use_layer_norm, num_hierarchies=num_hierarchies)

        new_model = copy.deepcopy(model)

    device = "cpu"

    h1_encoder_params = []
    h1_decoder_params = []

    named_params = list(new_model.named_parameters())

    index = 0

    for name, param in named_params:

        print("i: " + str(index) + " ", name, param.shape)
        index += 1

        if "h1" in name:
            if "encode" in name:
                h1_encoder_params += [param]
            elif "decode" in name:
                h1_decoder_params += [param]
            else:
                h1_decoder_params += [param]

    h_param_groups = [{'params': h1_encoder_params, 'g_name': 'h1_encoder_params'},
                      {'params': h1_decoder_params, 'g_name': 'h1_decoder_params'},
                      ]

    new_optimizer = optim.AdamW(h_param_groups, lr=Adam_learning_rate, weight_decay=0.0)

    set_param_grad_status(new_optimizer, ['h1_encoder_params', 'h1_decoder_params'], [])

    if bootstrap == "no":
        num_epochs = main_initial_num_epochs
    else:
        num_epochs = anneal_logistic_epochs(0, bootstrap=bootstrap)

    current_time = datetime.now().strftime('%b%d_%H-%M-%S')

    z_recon_stats_file = ('z_recon_stats/' + 'full_' + current_time + '.txt')

    if bootstrap != "no":

        z_seq_len = seq_len
        z_grid = initial_region_points(num_grid_points, z_size)
        z_grid = torch.from_numpy(z_grid).float().to(device)

        if small_steps:
            pos_to_sample = torch.cat([z_grid]*4)
        else:
            pos_to_sample = z_grid

        trajectory_actions, trajectory_actions_one_hot, _, trajectory_env_states = \
            new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                            argmax=True, env=env,
                                                            bootstrap_level=bootstrap, old_model=model)

    for compress_iteration in range(num_compression_iterations):

        print(colorize("compress_iteration: ", 'red'), compress_iteration)
        print("num_epochs: ", num_epochs)

        log_dir = os.path.join('LSL_new_runs', current_time + '_' +
                               anneal + '_num_epochs_' + str(num_epochs) +
                               "comp_iter_" + str(compress_iteration) + '_' + str(torch_manual_seed))

        print(log_dir)

        run_VAE_statsfile = ('LSL_run_VAE_stats/' + 'full_' +
                             anneal + '_num_epochs_' + str(num_epochs) +
                             str(torch_manual_seed) + "_comp_iter_" + str(compress_iteration) + current_time + '.txt')

        print(colorize("Writing stats to: " + run_VAE_statsfile, 'green'))
        f = open(run_VAE_statsfile, 'a')

        writer = SummaryWriter(log_dir)

        if small_steps:
            num_epochs = 3
            anneal = 'const'

        run_VAE(new_model, new_optimizer, trajectory_env_states, trajectory_actions, trajectory_actions_one_hot,
                num_epochs=num_epochs, clip_gn=clip_gn, clip_gv=clip_gv, batch_size=batch_size,
                kl_tolerance=kl_tolerance, anneal=anneal, logistic_scale=logistic_scale, target_beta=target_beta,
                floor=floor, log_interval=1, log_epoch_interval=1, f=f, writer=writer)

        iteration_model_file = ('saved_models_full/model_h1'
                                + anneal + 'comp_iter' + str(compress_iteration) +
                                str(torch_manual_seed) + current_time + '.pt')

        torch.save(new_model, iteration_model_file)

        if compress_iteration == num_compression_iterations - 1:

            final_model_file = ('saved_models_full/model_h1_final' +
                                anneal + 'comp_iter' + str(compress_iteration) +
                                str(torch_manual_seed) + given_time + '.pt')

            torch.save(new_model, final_model_file)

        if compress_iteration == 0 or compress_iteration == num_compression_iterations - 1:
            f.close()
            writer.close()

        print(colorize("-----Finished VAE iteration-----\n", 'red'))

        z_grid = initial_region_points(num_grid_points, z_size)
        z_grid = torch.from_numpy(z_grid).float().to(device)

        z_seq_len = seq_len
        primitive_trajectory_len = 5

        num_samples = z_grid.shape[0]
        new_model.eval()

        pos_to_sample = find_good_region(new_model, z_grid, z_seq_len, primitive_trajectory_len,
                                         num_samples=num_samples, env=env,
                                         bootstrap=bootstrap, old_model=model)

        pos_to_sample = pos_to_sample[:num_samples]

        if small_steps:
            pos_to_sample = torch.cat([pos_to_sample]*4)

        print(colorize("-----Finished finding sample region-----\n", 'red'))
        new_model.eval()

        if bootstrap == "no":

            trajectory_actions_argmax, trajectory_actions_one_hot_argmax, \
                trajectory_probs_argmax, trajectory_states_argmax = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    argmax=True, env=env)
        else:
            trajectory_actions_argmax, trajectory_actions_one_hot_argmax, \
                trajectory_probs_argmax, trajectory_states_argmax = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                              argmax=True, env=env,
                                                                              bootstrap_level=bootstrap,
                                                                              old_model=model)

        print(colorize("-----Finished argmax sampling-----\n", 'red'))

        start_time = time.time()

        if bootstrap == "no" and not small_steps:
            high_trajectory_actions_repulsion, high_trajectory_actions_one_hot_repulsion, \
            high_trajectory_probs_repulsion, high_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=7, high_repulsion_idx=9,
                                                                    argmax=False, env=env)

            mid_trajectory_actions_repulsion, mid_trajectory_actions_one_hot_repulsion, \
            mid_trajectory_probs_repulsion, mid_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=3, high_repulsion_idx=6,
                                                                    argmax=False, env=env)

            low_trajectory_actions_repulsion, low_trajectory_actions_one_hot_repulsion, \
            low_trajectory_probs_repulsion, low_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region(pos_to_sample, z_seq_len, primitive_trajectory_len,
                                                                    low_repulsion_idx=0, high_repulsion_idx=2,
                                                                    argmax=False, env=env)
        elif not small_steps:
            high_trajectory_actions_repulsion, high_trajectory_actions_one_hot_repulsion, \
            high_trajectory_probs_repulsion, high_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(
                                                            pos_to_sample, z_grid, z_seq_len,
                                                            low_repulsion_idx=7, high_repulsion_idx=9,
                                                            argmax=False, env=env,
                                                            bootstrap_level=bootstrap, old_model=model)

            mid_trajectory_actions_repulsion, mid_trajectory_actions_one_hot_repulsion, \
            mid_trajectory_probs_repulsion, mid_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                    low_repulsion_idx=3, high_repulsion_idx=6,
                                                                    argmax=False, env=env,
                                                                    bootstrap_level=bootstrap, old_model=model)

            low_trajectory_actions_repulsion, low_trajectory_actions_one_hot_repulsion, \
            low_trajectory_probs_repulsion, low_trajectory_states_repulsion = \
                new_model.argmax_or_repulsion_primitive_from_region_bootstrap(pos_to_sample, z_grid, z_seq_len,
                                                                    low_repulsion_idx=0, high_repulsion_idx=2,
                                                                    argmax=False, env=env,
                                                                    bootstrap_level=bootstrap, old_model=model)

        print("repulsion time: " + str(time.time() - start_time))
        print(colorize("-----Finished repulsion sampling-----\n", 'red'))

        num_trajectories = trajectory_actions_argmax.shape[0]

        trajectory_diagnostics(model=new_model, z_space_folder='z_space_trajectories', env=env,
                               epoch=compress_iteration,
                               z_grid=z_grid[:num_trajectories],
                               z_seq_len=z_seq_len, primitive_trajectory_len=5,
                               plot=True, show_plot=False,
                               z_recon_stats_file=z_recon_stats_file)

        if compress_iteration < main_num_logistic_compression_iterations and (argmax_only is False):

            num_epochs = anneal_logistic_epochs(compress_iteration, bootstrap="no")

            trajectory_actions = torch.cat([trajectory_actions_argmax,
                                            high_trajectory_actions_repulsion,
                                            mid_trajectory_actions_repulsion,
                                            low_trajectory_actions_repulsion])

            trajectory_actions_one_hot = torch.cat([trajectory_actions_one_hot_argmax,
                                                    high_trajectory_actions_one_hot_repulsion,
                                                    mid_trajectory_actions_one_hot_repulsion,
                                                    low_trajectory_actions_one_hot_repulsion])

            trajectory_env_states = torch.cat([trajectory_states_argmax,
                                               high_trajectory_states_repulsion,
                                               mid_trajectory_states_repulsion,
                                               low_trajectory_states_repulsion])
        else:

            num_epochs = main_num_const_epochs
            anneal = 'const'

            trajectory_actions = trajectory_actions_argmax
            trajectory_actions_one_hot = trajectory_actions_one_hot_argmax
            trajectory_env_states = trajectory_states_argmax

        if not small_steps:

            random_indices = np.arange(num_trajectories)
            np.random.shuffle(random_indices)

    print(colorize("-----Finished final VAE iteration-----\n", 'red'))
    new_model.eval()

    torch.save(new_model, model_file)
    new_model = torch.load(model_file)
    new_model.eval()

    return final_model_file

def anneal_logistic_epochs(compress_iteration, bootstrap="no"):

    if bootstrap == "no":
        if compress_iteration < 5:
            return 5
        if compress_iteration < 10:
            return 8
        if compress_iteration < 15:
            return 8

        return 8
    else:
        return 5