import PyQt4.QtGui as qg
import PyQt4.QtCore as qc

import maya.cmds as mc

import os
import xml.etree.ElementTree as ET

import cg_inventor.sys.lib.qt.base as base
import cg_inventor.sys.lib.qt.widgets.scene_check as scene_check

from cg_inventor.maya.apps.toolbox import page, settings

import cg_inventor.maya.utils.generic    as util_generic
import cg_inventor.maya.utils.transforms as util_transforms
import cg_inventor.maya.utils.history    as util_history
import cg_inventor.maya.utils.mesh       as util_mesh
import cg_inventor.maya.utils.texture    as util_texture

callbackPlus    = util_generic.callbackPlus
callbackSetup   = util_generic.callbackSetup
callbackCounter = util_generic.callbackCounter

setup  = scene_check.setup
check  = scene_check.check
select = scene_check.select
fix    = scene_check.fix

MODE_SEL = scene_check.SceneCheck.MODE_SEL
MODE_ALL = scene_check.SceneCheck.MODE_ALL

# ------------------------------------------------------------------------------------------------ #
#                                      MODELLING CHECKS                                            #
# ------------------------------------------------------------------------------------------------ #

TITLE    = 'Scene Checker'
SUBTITLE = 'Modelling'

class ModellingChecks(page.Page):
    def __init__(self):
        page.Page.__init__(self)
        self.title    = TITLE
        self.subtitle = SUBTITLE

        self._all_sets    = []
        self._all_spacers = []

        self.settings_widget = ModellingChecksSettings

        self.loadXML("scene_checks_modelling.xml")

    def _addSceneCheckSet(self, title):
        new_set = scene_check.SceneCheckSet(title)
        self._all_sets.append(new_set)
        self.addElement(new_set)

        self.connect(new_set, qc.SIGNAL(base.Base.STATUS_NORMAL),  self.statusNormal)
        self.connect(new_set, qc.SIGNAL(base.Base.STATUS_SUCCESS), self.statusSuccess)
        self.connect(new_set, qc.SIGNAL(base.Base.STATUS_ERROR),   self.statusError)

        return new_set

    def clear(self):
        for check_set in self._all_sets:
            self.removeElement(check_set)
            check_set.deleteLater()

        for spacer in self._all_spacers:
            self.removeElement(spacer)

        self._all_sets    = []
        self._all_spacers = []

    def loadXML(self, xml_file):
        # get toolbox xml path
        #
        import cg_inventor.maya.apps.toolbox as toolbox
        xml_path = os.path.join(os.path.dirname(toolbox.__file__), 'xml', xml_file)

        # get the root element
        #
        tree = ET.parse(xml_path)
        root = tree.getroot()

        for check_set in root:
            set_title = check_set.attrib['title']
            new_set = self._addSceneCheckSet(set_title)

            for check in check_set:
                function_set = str(check.text).strip()
                new_check = scene_check.SceneCheck()
                new_check.loadFromLibrary(function_set)

                new_set.addSceneCheck(new_check)

            spacer = qg.QSpacerItem(20,20,qg.QSizePolicy.Fixed)
            self.addElement(spacer)
            self._all_spacers.append(spacer)

# ------------------------------------------------------------------------------------------------ #
#                                      MODELLING CHECKS                                            #
# ------------------------------------------------------------------------------------------------ #

class ModellingChecksSettings(settings.Settings):
    def __init__(self):
        settings.Settings.__init__(self)
        #self.setMinimumWidth(600)
        self.setFixedWidth(700)

        main_layout = qg.QHBoxLayout()
        main_layout.setContentsMargins(5,0,5,5)
        main_layout.setSpacing(5)
        self.addElement(main_layout)

        group_widget      = qg.QWidget()
        group_widget.setLayout(qg.QVBoxLayout())
        group_widget.layout().setContentsMargins(0,0,0,0)
        group_widget.layout().setSpacing(0)
        check_widget      = qg.QWidget()
        check_widget.setLayout(qg.QVBoxLayout())
        check_widget.layout().setContentsMargins(0,0,0,0)
        check_widget.layout().setSpacing(0)
        button_widget     = qg.QWidget()
        button_widget.setLayout(qg.QVBoxLayout())
        all_checks_widget = qg.QWidget()
        all_checks_widget.setLayout(qg.QVBoxLayout())
        all_checks_widget.layout().setContentsMargins(0,0,0,0)
        all_checks_widget.layout().setSpacing(0)
        description_widget = qg.QWidget()
        description_widget.setLayout(qg.QVBoxLayout())
        description_widget.layout().setContentsMargins(0,0,0,0)
        description_widget.layout().setSpacing(0)

        button_widget.setFixedWidth(30)

        main_layout.addWidget(group_widget)
        main_layout.addWidget(check_widget)
        main_layout.addWidget(button_widget)
        main_layout.addWidget(all_checks_widget)
        main_layout.addWidget(description_widget)

        title_font = qg.QFont()
        title_font.setBold(True)

        group_label        = qg.QLabel('GROUPS')
        group_layout       = qg.QHBoxLayout()
        checks_label       = qg.QLabel('CHECKS')
        checks_layout      = qg.QHBoxLayout()
        all_checks_label   = qg.QLabel('ALL CHECKS')
        all_checks_layout  = qg.QHBoxLayout()
        description_label  = qg.QLabel('DESCRIPTION')
        description_layout = qg.QHBoxLayout()

        group_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))
        group_layout.addWidget(group_label)
        group_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))

        checks_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))
        checks_layout.addWidget(checks_label)
        checks_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))

        all_checks_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))
        all_checks_layout.addWidget(all_checks_label)
        all_checks_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))

        description_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))
        description_layout.addWidget(description_label)
        description_layout.addSpacerItem(qg.QSpacerItem(5, 5, qg.QSizePolicy.Expanding))

        group_label.setFont(title_font)
        checks_label.setFont(title_font)
        all_checks_label.setFont(title_font)
        description_label.setFont(title_font)

        group_label.setFixedHeight(18)
        checks_label.setFixedHeight(18)
        all_checks_label.setFixedHeight(18)
        description_label.setFixedHeight(18)

        self.group_list      = qg.QListWidget()
        self.check_list      = qg.QListWidget()
        self.all_checks_list = qg.QListWidget()
        self.description_te  = qg.QTextEdit()

        group_widget.layout().addLayout(group_layout)
        group_widget.layout().addWidget(self.group_list)
        check_widget.layout().addLayout(checks_layout)
        check_widget.layout().addWidget(self.check_list)
        all_checks_widget.layout().addLayout(all_checks_layout)
        all_checks_widget.layout().addWidget(self.all_checks_list)
        description_widget.layout().addLayout(description_layout)
        description_widget.layout().addWidget(self.description_te)

        self.all_checks_list.itemClicked.connect(self.showDescription)

        self._check_descriptions = {}

        self.populateAllChecks()

    def clearAllChecks(self):
        pass

    def populateAllChecks(self):
        self.clearAllChecks()

        all_names = []
        for check_name, setup_func in scene_check.SceneCheckLibrary.setup.items():
            title, _, description = setup_func()
            self._check_descriptions[title] = description
            all_names.append(title)

        for name in sorted(all_names):
            new_item = qg.QListWidgetItem()
            new_item.setText(name)
            self.all_checks_list.addItem(new_item)

    def showDescription(self, check_item):
        description = self._check_descriptions[str(check_item.text())]

        self.description_te.setText(description)

# ------------------------------------------------------------------------------------------------ #
#                                           CHECKS                                                 #
# ------------------------------------------------------------------------------------------------ #

@setup
def freezeTransforms():
    title       = 'Unfrozen Transforms'
    message     = 'Found $NUM Unfrozen Transform(s)'
    description = 'Finds all transforms with unfrozen values.'
    return title, message, description

@check
def freezeTransforms(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            transforms = mc.ls(sl=True)
        else:
            transforms = util_transforms.getAllTransforms()

        if not len(transforms): return

        return util_transforms.getUnfrozenTransforms(transforms, callback=callback)

    except util_transforms.TransformError as e:
        raise scene_check.SceneCheckError(e)

@select
def freezeTransforms(unfrozen_transforms):
    mc.select(mc.ls(unfrozen_transforms, type='transform'), r=True)

@fix
def freezeTransforms(transforms, callback=None):
    try:
        util_transforms.freezeTransforms(transforms, filter=True, callback=callback)
    except util_transforms.TransformError as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def nonDeletedHistory():
    title       = 'Non-Deleted History'
    message     = 'Found $NUM Node(s) with History'
    description = 'Find any nodes with history.'
    return title, message, description

@check
def nonDeletedHistory(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            transforms = mc.ls(sl=True)
        else:
            transforms = util_transforms.getAllTransforms()

        if not len(transforms): return

        return util_history.getNodesWithHistory(transforms, callback=callback)
    except util_history.HistoryError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def nonDeletedHistory(history_nodes):
    mc.select(mc.ls(history_nodes), r=True)

@fix
def nonDeletedHistory(history_nodes):
    try:
        util_history.deleteHistory(history_nodes)
    except util_history.HistoryError as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def centeredPivots():
    title       = 'Non Centered Pivots'
    message     = 'Found $NUM Non Centered Pivot(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def centeredPivots(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        if callback: counter, increment = callbackSetup(len(meshes), callback_range=30)

        transforms = set([])
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            if util_mesh._isMeshEmpty(mesh): continue
            transforms.add(util_mesh.getTransformFromShape(mesh))

        if callback: counter, increment = callbackSetup(len(transforms), callback_range=70)

        not_centred = []
        for transform in transforms:
            if callback: callback(30 + callbackCounter(counter, increment))

            if not util_transforms.isPivotCentered(transform):
                not_centred.append(transform)

        return not_centred

    except (util_transforms.TransformError, util_mesh.MeshError) as e:
        raise scene_check.SceneCheckError(str(e))

@select
def centeredPivots(transforms):
    mc.select(transforms, r=True)

@fix
def centeredPivots(transforms, callback=None):
    try:
        mc.xform(transforms, cp=True)
    except Exception as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def unknownNodes():
    title       = "'Unknown' Nodes"
    message     = "Found $NUM 'Unknown' Node(s)"
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def unknownNodes(selection_mode):
    if selection_mode == MODE_SEL:
        return mc.ls(sl=True, type='unknown')
    else:
        return mc.ls(type='unknown')

@select
def unknownNodes(nodes):
    mc.select(mc.ls(nodes), r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def triangles():
    title       = 'Triangles'
    message     = 'Found $NUM Triangle(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def triangles(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        return util_mesh.checkMesh(meshes, flags=util_mesh.TRIANGLES, callback=callback)

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def triangles(triangles):
    mc.select(triangles, r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def nsided():
    title   = 'N-Sided Faces'
    message     = 'Found $NUM N-Sided Face(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def nsided(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        return util_mesh.checkMesh(meshes, flags=util_mesh.NSIDED, callback=callback)

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def nsided(faces):
    mc.select(faces, r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def holed():
    title       = 'Holes'
    message     = 'Found $NUM Hole(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def holed(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        if callback: counter, increment = callbackSetup(len(meshes), callback_range=30)

        holes = []
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            if util_mesh._isMeshEmpty(mesh): continue
            mesh_holes = util_mesh.getMeshHoles(mesh)
            holes.extend([(mesh, mesh_hole) for mesh_hole in mesh_holes])

        return holes

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def holed(holes):
    edges = []
    for mesh, edge_indices in holes:
        edges.extend(['%s.e[%s]' %(mesh, edge_index) for edge_index in edge_indices])
    mc.select(mc.ls(edges), r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def zeroLengthEdges():
    title       = 'Zero Length Edges'
    message     = 'Found $NUM Zero Length Edge(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def zeroLengthEdges(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        return util_mesh.checkMesh(meshes, flags=util_mesh.ZERO_LENGTH_EDGES, callback=callback)

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def zeroLengthEdges(edges):
    mc.select(mc.ls(edges), r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def zeroAreaFaces():
    title       = 'Zero Area Faces'
    message     = 'Found $NUM Zero Area Face(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def zeroAreaFaces(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        if not len(meshes): return

        return util_mesh.checkMesh(meshes, flags=util_mesh.ZERO_AREA_FACES, callback=callback)

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def zeroAreaFaces(faces):
    mc.select(mc.ls(faces), r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def lockedNormals():
    title       = 'Locked Normals'
    message     = 'Found $NUM Locked Normal(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def lockedNormals(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        num_meshes = len(meshes)
        if not num_meshes: return

        if callback: counter, increment = callbackSetup(num_meshes, callback_range=100)

        locked_normals = []
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            if util_mesh._isMeshEmpty(mesh): continue
            indices = util_mesh.getLockedVertexNormals(mesh)
            vertices = ['%s.vtx[%s]' %(mesh, index) for index in indices]
            locked_normals.extend(vertices)

        return locked_normals

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def lockedNormals(normals):
    mc.select(mc.ls(normals), r=True)

@fix
def lockedNormals(normals):
    try:
        mc.polyNormalPerVertex(normals, ufn=True)
    except Exception as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def hardEdges():
    title       = 'Hard Edges'
    message     = 'Found $NUM Hard Edge(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def hardEdges(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        num_meshes = len(meshes)
        if not num_meshes: return

        if callback: counter, increment = callbackSetup(num_meshes, callback_range=100)

        hard_edges = []
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            if util_mesh._isMeshEmpty(mesh): continue
            hard_edges.extend([(mesh, edge) for edge in util_mesh.getHardEdges(mesh)])

        return hard_edges

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def hardEdges(edges):
    edge_strings = []
    for mesh, edge in edges:
        edge_strings.append('%s.e[%s]' %(mesh, edge))
    mc.select(edge_strings, r=True)

@fix
def hardEdges(edges, callback=None):
    try:
        all_hard_edges = {}

        if callback:
            counter, increment = callbackSetup(len(edges), callback_range=20)

        for mesh, edge in edges:
            if callback: callback(callbackCounter(counter, increment))

            try:
                hard_edges = all_hard_edges[mesh]
            except KeyError:
                hard_edges = all_hard_edges[mesh] = []
            hard_edges.append('%s.e[%s]' %(mesh, edge))

        if callback:
            counter, increment = callbackSetup(len(all_hard_edges.keys()), callback_range=80)

        for mesh, hard_edges in all_hard_edges.items():
            if callback: callback(20 + callbackCounter(counter, increment))

            history = bool(mc.listHistory(mesh, il=2, pdo=True))
            mc.polySoftEdge(*hard_edges, a=180, ch=history)

    except Exception as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def unconnectedIntermediates():
    title       = 'Unconnected Intermediates'
    message     = 'Found $NUM Unconnected Intermediate(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def unconnectedIntermediates(selection_mode):
    return util_mesh.getUnconnectedIntermediates()

@select
def unconnectedIntermediates(intermediates):
    mc.select(mc.ls(intermediates), r=True)

@fix
def unconnectedIntermediates(intermediates):
    intermediates = mc.ls(intermediates)
    if not intermediates: return

    mc.delete(intermediates)

#--------------------------------------------------------------------------------------------------#

@setup
def emptyMeshes():
    title       = 'Empty Meshes'
    message     = 'Found $NUM Empty Mesh(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def emptyMeshes(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getEmptyMeshes()

        num_meshes = len(meshes)
        if not num_meshes: return

        if callback: counter, increment = callbackSetup(num_meshes, callback_range=100)

        empty_meshes = []
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            if util_mesh._isMeshEmpty(mesh):
                empty_meshes.append(mesh)

        return meshes

    except util_mesh.MeshError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def emptyMeshes(meshes):
    mc.select(mc.ls(meshes), r=True)

@fix
def emptyMeshes(meshes):
    meshes = mc.ls(meshes)
    if not meshes: return

    try:
        mc.delete(meshes)
    except Exception as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def unusedShaders():
    title       = 'Unused Shaders'
    message     = 'Found $NUM Unused Shader(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def unusedShaders(selection_mode):
    try:
        if selection_mode == MODE_SEL:
            materials = mc.ls(sl=True, mat=True)

            unused_materials = []
            for material in materials:
                if util_texture._getMaterialMembers(material) is None:
                    unused_materials.append(material)

        else:
            unused_materials = util_texture.getAllUnusedMaterials()

        return unused_materials

    except util_texture.TextureError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def unusedShaders(materials):
    mc.select(materials, r=True)

@fix
def unusedShaders(materials, callback=None):
    try:
        if callback: counter, increment = callbackSetup(len(materials), callback_range=100)

        for material in materials:
            if callback: callback(callbackCounter(counter, increment))

            util_texture.deleteMaterial(material)

    except util_texture.TextureError as e:
        raise scene_check.SceneCheckError(str(e))

#--------------------------------------------------------------------------------------------------#

@setup
def uvsOutOfBounds():
    title       = 'UVs Out of Bounds'
    message     = 'Found $NUM UV(s) Out of Bounds'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def uvsOutOfBounds(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        num_meshes = len(meshes)
        if not num_meshes: return

        if callback: counter, increment = callbackSetup(num_meshes, callback_range=100)

        all_uvs_oob = []
        for mesh in meshes:
            if callback: callback(callbackCounter(counter, increment))

            uvs_oob = util_texture.getAllUVsOutOfBounds(mesh)
            for uv_index in uvs_oob:
                all_uvs_oob.append('%s.map[%s]' %(mesh, uv_index))

        return all_uvs_oob

    except util_texture.TextureError as e:
        raise scene_check.SceneCheckError(str(e))

@select
def uvsOutOfBounds(uvs):
    mc.select(uvs, r=True)

#--------------------------------------------------------------------------------------------------#

@setup
def fivePointers():
    title       = 'Five+ Pointers'
    message     = 'Found $NUM Five+ Pointer(s)'
    description = 'Find any geometry with uncentered pivots.'
    return title, message, description

@check
def fivePointers(selection_mode, callback=None):
    try:
        if selection_mode == MODE_SEL:
            meshes = util_mesh.getMeshesFromSelection(mc.ls(sl=True))
        else:
            meshes = util_mesh.getAllMeshes()

        five_pointers = []
        for mesh in meshes:
            vertices = util_mesh.getFiveEdgeVerts(mesh)
            for vertex_index in vertices:
                five_pointers.append('%s.vtx[%s]' %(mesh, vertex_index))

        return five_pointers

    except Exception as e:
        raise scene_check.SceneCheckError(str(e))

@select
def fivePointers(verts):
    mc.select(verts, r=True)


#!/usr/bin/python3
import datetime
from datetime import datetime, timedelta
import argparse
import sys
from lib.datacalc import datacalculator
from lib.checktotal import videocounter
from lib.youtubevideosearch import *
from lib.downloadcomments import *

#### VARIABLES #####

pAfter=None
pBefore=None
ptoken=""
stalkerlimit=10000

##### ARGS DATA FORMAT AND EXISTENCE VALIDATION ######

def validate_date(date):
    try:
        return datetime.strptime(date, "%Y-%m-%d").date()
    except ValueError:
        msg="Not a valid date: '{0}'".format(date)
        raise argparse.ArgumentTypeError(msg)
def validate_vorder(order):
    orders=["date","rating","relevance","title","videoCount","viewCount"]
    if order in orders:
        return order
    else:
        msg=ValueError("Invalid video order option.")
        raise argparse.ArgumentTypeError(msg)
def validate_corder(order):
    orders=["time","relevance"]
    if order in orders:
        return order
    else:
        msg=ValueError("Invalid comment order option.")
        raise argparse.ArgumentTypeError(msg)

if __name__ == '__main__':

    ########### ARGPARSE #############
    parser = argparse.ArgumentParser(prog='Youtube Stalker v.0.1', usage='%(prog)s [options]',description='Youtube Stalker is a simple python tool to help you monitor youtube videos comments about a topic of your interest and generate a output that can be used for data analytics and intelligence.',epilog='Usage example: youstalker.py -a 12-01-2018 -b 12-31-2018 -q "cats"')
    parser.add_argument('-q', '--query',help='Keyword - video search argument',action='store', required=True)
    parser.add_argument('-a', '--after', help='Published after - video search argument. Format YYYY-MM-DD', action="store", type=validate_date)
    parser.add_argument('-b', '--before',help='Published before - video search argument. Format YYYY-MM-DD', action="store", type=validate_date)
    parser.add_argument('-l','--video-limit', help='Video results limit. (default: 500). Maximum = {0}'.format(stalkerlimit), type=int, default='500')
    parser.add_argument('-o','--video-order',help='Possible orders: date, rating, relevance, title, videoCount, viewCount. (default: viewCount)',default='viewCount', action="store",type=validate_vorder)
    parser.add_argument('-L','--comment-limit', help='Comments limit for each video.(default: all comments)',type=int)
    parser.add_argument('-O','--comment-order', help='Possible orders: time, relevance. (default: time)', default='time',type=validate_corder)
    args = parser.parse_args()

    ########### ARGS VALUES VALIDATION ##########
    if args.video_limit > stalkerlimit:
        print("Stalker's limit is {0} videos for video search. Quitting...".format(stalkerlimit))
        sys.exit(1)
    if args.after and args.before:
        diff= args.after - args.before
        if diff.days > 0:
            print("--before and --after range is invalid!")
            sys.exit(1)
    if (args.video_limit > 500 and (args.after==None or args.before==None)):
        print("Youtube v3 API has a embeded video search limit of 500. \nIf you provide publishedAfter and publishedBefore values \nYoustalker will calculate the time range and split your \nquery into small queries to bypass this restriction and bring you 10000 videos results max.")
        sys.exit(1)
    if args.after:
        pAfter=(str(args.after)+"T00:00:01Z")
    if args.before:
        pBefore=(str(args.before)+"T23:59:59Z")

    ######### STALKER USE CASES ############

    if ((pAfter and pBefore) != None and args.limit > 500):
        print("Wait... I'm working on your request...")
        videototal=videocounter(args.query,pAfter,pBefore,ptoken,args.video_order)
        days=datacalculator(pAfter,pBefore,videototal)
        pega_comments(videos,args.comment_order,args.comment_limit)
        videos=youtube_searchdate(args.query,pAfter,pBefore,ptoken,args.video_order,days,args.video_limit,videototal)
    else:
        print("Wait... I'm working on your request...")
        videos=youtube_search(args.query,pAfter,pBefore,ptoken,args.video_order,args.video_limit)
        pega_comments(videos,args.comment_order,args.comment_limit)


import time
import io

def parseFastaFile1(infile, db):
    with open(infile, 'r') as inFile:
        for line in inFile:
            if line[0] == ">":
                ident, desc = line[1:].strip().split(maxsplit=1)
                db.append({'id': ident, 'description': desc, 'sequence': "", 'raw': line})
            else:
                db[-1]['sequence'] += line.rstrip()
                db[-1]['raw'] += line

def parseFastaFile2(infile, db):
    with open(infile, "rb") as inFile:
        for line in inFile:
            if line[0] == 62:  # bin token for ">"
                ident, desc = line[1:].strip().split(maxsplit=1)
                db.append({'id': ident,
                           'description': desc,
                           'sequence': bytearray(),
                           'raw': bytearray(line)})
            else:
                db[-1]['sequence'] += (line.rstrip())
                db[-1]['raw'] += (line)

# version using StringIO
def parseFastaFile3(infile, db):
    with open(infile, "r") as inFile:
        for line in inFile:
            if line[0] == ">":
                ident, desc = line[1:].strip().split(maxsplit=1)
                db.append({'id': ident,
                           'description': desc,
                           'sequence': io.StringIO(),
                           'raw': io.StringIO(line)})
            else:
                db[-1]['sequence'].write(line.rstrip())
                db[-1]['raw'].write(line)

# version using BytesIO
def parseFastaFile4(infile, db):
    with open(infile, "rb") as inFile:
        for line in inFile:
            if line[0] == 62:  # bin token for ">"
                ident, desc = line[1:].strip().split(maxsplit=1)
                db.append({'id': ident,
                           'description': desc,
                           'sequence': io.BytesIO(),
                           'raw': io.BytesIO(line)})
            else:
                db[-1]['sequence'].write(line.rstrip())
                db[-1]['raw'].write(line)

# Stop time of function calls
def stop_time(func, *args, **kargs):
    start = time.time()

    result = func(*args, **kargs)

    end = time.time()
    print("Function {} takes {:.3} seconds".format(func.__name__, end - start))
    return result

d = []

res1 = stop_time(parseFastaFile1, "../examples/long.fasta", db=d)
res2 = stop_time(parseFastaFile2, "../examples/long.fasta", db=d)
res3 = stop_time(parseFastaFile3, "../examples/long.fasta", db=d)
res4 = stop_time(parseFastaFile4, "../examples/long.fasta", db=d)

# Compare both sequences
print(d[0]['sequence'] == d[1]['sequence'].decode())
print(d[0]['sequence'] == d[2]['sequence'].getvalue())
print(d[0]['sequence'] == d[3]['sequence'].getvalue().decode())


import os
import re
import socket

import trello
from hugin.flowcells import Flowcell
from flowcell_status import FlowcellStatus, FC_STATUSES

FC_NAME_RE = r'(\d{6})_([ST-]*\w+\d+)_\d+_([AB]?)([A-Z0-9\-]+)'

COLORS = [
    'red',
    'blue',
    'green',
    'yellow',
    'orange',
    'purple',
    'lime',
    'pink',
    'sky',
    'black'
]

class FlowcellMonitor(object):
    def __init__(self, config):
        self._config = config
        # initialize None values for @property functions
        self._trello_board = None
        self._data_folders = None
        self._trello_cards = None
        self._trello_lists = None

    @property
    def config(self):
        return self._config

    @property
    def trello_board(self):
        if not self._trello_board:
            if not self.config.get('trello'):
                raise RuntimeError("'trello' must be in config file")

            config = self.config.get('trello')
            # todo check if board exist

            api_key = config.get('api_key')
            token = config.get('token')
            api_secret = config.get('api_secret')
            client = trello.TrelloClient(api_key=api_key, token=token, api_secret=api_secret)
            board_id = config.get('board_id')
            self._trello_board  = client.get_board(board_id)

        return self._trello_board

    @property
    def data_folders(self):
        if not self._data_folders:
            self._data_folders = self.config.get('data_folders')
            if self._data_folders is None:
                raise RuntimeError("'data_folders' must be in config file")
        return self._data_folders

    @property
    def trello_cards(self):
        if self._trello_cards is None:
            self._trello_cards = self.trello_board.all_cards()
        return self._trello_cards

    @property
    def trello_lists(self):
        if self._trello_lists is None:
            self._trello_lists = self.trello_board.all_lists()
        return self._trello_lists

    def update_trello_board(self):
        for data_folder in self.data_folders:
            self._check_running_flowcells(data_folder)
            self._check_nosync_flowcells(data_folder)
            # move deleted flowcells to the archive list
            self._check_archived_flowcells(data_folder)

    def _check_running_flowcells(self, data_folder):
        # go through subfolders
        subfolders = filter(os.path.isdir, [os.path.join(data_folder, fc_path) for fc_path in os.listdir(data_folder)])
        for flowcell_path in subfolders:
            # skip non-flowcell folders
            if not re.match(FC_NAME_RE, os.path.basename(flowcell_path)):
                continue

            status = FlowcellStatus(flowcell_path)
            # depending on the type, return instance of related class (hiseq, hiseqx, miseq, etc)
            flowcell = Flowcell.init_flowcell(status)
            if flowcell.check_status():
                # todo: add comment
                # todo: if comment has been added
                pass
            # update flowcell on trello board
            self._update_card(flowcell)

    def _check_nosync_flowcells(self, data_folder):
        # check nosync folder
        nosync_folder = os.path.join(data_folder, 'nosync')
        if os.path.exists(nosync_folder):
            # move flowcell to nosync list
            for nosync_flowcell in os.listdir(nosync_folder):
                flowcell_path = os.path.join(nosync_folder, nosync_flowcell)
                # skip non-flowcell folders
                if not re.match(FC_NAME_RE, os.path.basename(flowcell_path)):
                    continue
                card = self._get_card_by_name(nosync_flowcell)
                # if the card is not on Trello board, create it
                if card is None:
                    status = FlowcellStatus(flowcell_path)
                    flowcell = Flowcell.init_flowcell(status)
                    self._update_card(flowcell)
                else:
                    nosync_list = self._get_list_by_name(FC_STATUSES['NOSYNC'])
                    card.change_list(nosync_list.id)

    def _check_archived_flowcells(self, data_folder):
        # if nosync folder exists
        if os.path.exists(os.path.join(data_folder, FC_STATUSES['NOSYNC']).lower()):
            # get cards from the nosync list
            for card in self._get_cards_by_list(FC_STATUSES['NOSYNC']):
                localhost = socket.gethostname()
                # if the flowcell belongs to the server
                if localhost in card.description:
                    # check if the flowcell has been deleted from the nosync folder
                    if card.name not in os.listdir(os.path.join(data_folder, FC_STATUSES['NOSYNC'].lower())):
                        archived_list = self._get_list_by_name(FC_STATUSES['ARCHIVED'])
                        card.change_list(archived_list.id)

    def _update_card(self, flowcell):
        # todo: beautify the method
        trello_card = self._get_trello_card(flowcell) # None
        flowcell_list = self._get_list_by_name(flowcell.trello_list)

        # if not card on trello board
        if trello_card is None:
            return self._create_card(flowcell)
        else:
            # skip aborted list
            if flowcell.trello_list == FC_STATUSES['ABORTED']:
                return trello_card
            # if card is in the wrong list
            if trello_card.list_id != flowcell_list.id:
                # move card
                trello_card.change_list(flowcell_list.id)

            # if card is in the right list
            else:
                # todo: checkstatus -> taking too long?
                return trello_card

            # update due_time
            trello_card.set_due(flowcell.due_time)
            if flowcell.trello_list == FC_STATUSES['CHECKSTATUS']:
                trello_card.comment(flowcell.status.warning)
            return trello_card

    def _create_card(self, flowcell):
        trello_list = self._get_list_by_name(flowcell.trello_list)
        if not trello_list:
            raise RuntimeError('List {} cannot be found in TrelloBoard {}'.format(flowcell.status, self.trello_board))

        trello_card = trello_list.add_card(name=flowcell.full_name, desc=flowcell.get_formatted_description())
        if flowcell.trello_list == FC_STATUSES['CHECKSTATUS']:
            trello_card.comment(flowcell.status.warning)
        trello_card.set_due(flowcell.due_time)
        self._add_label(trello_card, flowcell)

    def _add_label(self, card, flowcell):
        server = flowcell.server
        label = self._get_label_by_name(server)
        if label is None:
            color = self._get_next_color()
            label = self.trello_board.add_label(name=server, color=color)
        if label.id not in [label.id for label in card.labels]:
            card.add_label(label)

    def _get_label_by_name(self, name):
        labels = self.trello_board.get_labels()
        for label in labels:
            if label.name == name:
                return label
        return None

    def _get_list_by_name(self, list_name):
        for item in self.trello_lists:
            if item.name == list_name:
                return item
        return None

    def _get_cards_by_list(self, list_name):
        trello_list = self._get_list_by_name(list_name)
        result = []
        for card in self.trello_cards:
            if card.list_id == trello_list.id:
                result.append(card)
        return result

    def _get_card_by_name(self, card_name):
        for card in self.trello_cards:
            if card.name == card_name:
                return card

    def _get_trello_card(self, flowcell):
        for card in self.trello_cards:
            if flowcell.full_name == card.name:
                return card
        return None

    def _get_next_color(self):
        labels = self.trello_board.get_labels()
        colors = [label.color for label in labels] if labels else []
        # if all colors are used take the first one
        if colors == COLORS:
            return COLORS[0]

        # if not all the colors are used, take the one which is not used
        elif set(colors) != set(COLORS):
            for color in COLORS:
                if color not in colors:
                    return color

        else: # set(colors) == set(COLORS):
            # otherwise take the color which is used the least
            color_groups = {} # how many times each color has been used already
            for color in COLORS:
                color_groups[color] = colors.count(color)

            for color, count in color_groups:
                if count == min(color_groups.values()):
                    return color


import os.path as op
import glob
import numpy as np
import pandas as pd
import subprocess
import spacy
from wordfreq import zipf_frequency as word_frequency
import subprocess
try:
    spacy.load('nl_core_news_sm')
except:
    process = subprocess.Popen('python -m spacy download nl_core_news_sm'.split(), stdout=subprocess.PIPE)
    output, error = process.communicate()

# Setup paths and file names
this_path = '/' + op.join(*op.realpath(__file__).split('/')[:-1])

def setup_datapath():
    fname_datapath = '/' + this_path + '/data_path.txt'

    if not op.isfile(fname_datapath):
        data_path = input("Enter data path (or create a data_path.txt):")
        if data_path[-1] != '/':
            data_path += '/'
        with open(fname_datapath, 'w') as f:
            f.write(data_path)
    with open(fname_datapath, 'r') as f:
        data_path = f.read().strip('\n')
    return data_path

def setup_logfiles():
    fname_logfiles = '/' + this_path + '/log_files.csv'
    data_path = setup_datapath()
    if not op.isfile(fname_logfiles):
        tasks = dict(visual='Vis', auditory='Aud')
        log_files = list()

        for task in ('visual', 'auditory'):

            log_path = data_path + 'sourcedata/meg_task/'
            files = glob.glob(log_path + "*-MEG-MOUS-%s*.log" % tasks[task])

            for file in np.sort(files):
                subject = file.split('-')[0].split('/')[-1]
                log_files.append(dict(
                    subject=int(subject[1:]),
                    task=task,
                    log_id=int(file.split('-')[1]),
                    log_file=op.join('sourcedata', 'meg_task',
                                     file.split('/')[-1]),
                    meg_file=op.join('sub-' + subject, 'meg',
                                     'sub-%s_task-%s_meg.ds' % (subject, task))
                    )
                )
        log_files = pd.DataFrame(log_files)
        # Remove corrupted log
        log_files = log_files.loc[(log_files.subject != 1006) &
                                  (log_files.subject != 1017)]
        log_files.to_csv(fname_logfiles)
    return pd.read_csv(fname_logfiles)

def setup_stimuli():
    fname_stimuli = '/' + this_path + '/stimuli.csv'
    data_path = setup_datapath()
    if not op.isfile(fname_stimuli):

        source = op.join(data_path, 'stimuli', 'stimuli.txt')

        with open(source, 'r') as f:
            stimuli = f.read()
        while '  ' in stimuli:
            stimuli.replace('  ', ' ')

        # clean up
        stimuli = stimuli.split('\n')[:-1]
        stim_id = [int(s.split(' ')[0]) for s in stimuli]
        sequences = [' '.join(s.split(' ')[1:]) for s in stimuli]
        stimuli = pd.DataFrame([dict(index=idx, sequence=seq)
                                for idx, seq in zip(stim_id, sequences)])
        stimuli.to_csv(fname_stimuli, index=False)
    return pd.read_csv(fname_stimuli, index_col='index')

def setup_morphemes():
    # download
    commands = ('polyglot download morph2.en morph2.ar',
                'polyglot download morph2.nl')
    for command in commands:
        subprocess.Popen(command.split(), stdout=subprocess.PIPE)

data_path = setup_datapath()
log_files = setup_logfiles()
stimuli = setup_stimuli()

def _parse_log(log_fname):
    with open(log_fname, 'r') as f:
        text = f.read()

    # Fixes broken inputs
    text = text.replace('.\n', '.')

    # file is made of two blocks
    block1, block2 = text.split('\n\n\n')

    # read first header
    header1 = block1.replace(' ', '_').split('\n')[3].split('\t')
    header1[6] = 'time_uncertainty'
    header1[8] = 'duration_uncertainty'

    # read first data
    df1 = pd.DataFrame([s.split('\t') for s in block1.split('\n')][5:],
                       columns=header1)
    # the two dataframe are only synced on certains rows
    common_samples = ('Picture', 'Sound', 'Nothing')
    sel = df1['Event_Type'].apply(lambda x: x in common_samples)
    index = df1.loc[sel].index

    # read second header
    header2 = block2.replace(' ', '_').split('\n')[0].split('\t')
    header2[7] = 'time_uncertainty'
    header2[9] = 'duration_uncertainty'

    # read second data
    df2 = pd.DataFrame([s.split('\t') for s in block2.split('\n')[2:-1]],
                       columns=header2, index=index)

    # remove duplicate
    duplicates = np.intersect1d(df1.keys(), df2.keys())
    for key in duplicates:
        assert (df1.loc[index, key] == df2[key].fillna('')).all()
        df2.pop(key)

    log = pd.concat((df1, df2), axis=1)
    return log

def _clean_log(log):
    # Relabel condition: only applies to sample where condition changes
    translate = dict(
        ZINNEN='sentence',
        WOORDEN='word_list',
        FIX='fix',
        QUESTION='question',
        Response='response',
        ISI='isi',
        blank='blank',
    )
    for key, value in translate.items():
        sel = log.Code.astype(str).str.contains(key)
        log.loc[sel, 'condition'] = value
    log.loc[log.Code == '', 'condition'] = 'blank'

    # Annotate sequence idx and extend context to all trials
    start = 0
    block = 0
    context = 'init'
    log['new_context'] = False
    query = 'condition in ("word_list", "sentence")'
    for idx, row in log.query(query).iterrows():
        log.loc[start:idx, 'context'] = context
        log.loc[start:idx, 'block'] = block
        log.loc[idx, 'new_context'] = True
        context = row.condition
        block += 1
        start = idx
    log.loc[start:, 'context'] = context
    log.loc[start:, 'block'] = block

    # Format time
    log['time'] = 0
    idx = log.Time.str.isnumeric() == True  # noqa
    log.loc[idx, 'time'] = log.loc[idx, 'Time'].astype(float) / 1e4

    # Extract individual word
    log.loc[log.condition.isna(), 'condition'] = 'word'
    idx = log.condition == 'word'
    words = log.Code.str.strip('0123456789 ')
    log.loc[idx, 'word'] = words.loc[idx]
    sel = log.query('word=="" and condition=="word"').index
    log.loc[sel, 'word'] = pd.np.nan
    log.loc[log.word.isna() & (log.condition == "word"), 'condition'] = 'blank'
    return log

def _add_stim_id(log, verbose):
    # Find beginning of each sequence (word list or sentence)
    start = 0
    sequence_pos = -1
    for idx, row in log.query('condition == "fix"').iterrows():
        if sequence_pos >= 0:
            log.loc[start:idx, 'sequence_pos'] = sequence_pos
        sequence_pos += 1
        start = idx
    log.loc[start:, 'sequence_pos'] = sequence_pos

    # Find corresponding stimulus id
    stim_id = 0
    lower30 = lambda s: s[:30].lower()  # noqa
    stimuli['first_30_chars'] = stimuli.sequence.apply(lower30)
    sel = slice(0, 0)
    for pos, row in log.groupby('sequence_pos'):
        if pos == -1:
            continue

        # select words in this sequence
        sel = row.condition == "word"
        if not sum(sel):
            continue

        # match with stimuli
        first_30_chars = ' '.join(row.loc[sel, 'word'])[:30].lower()  # noqa
        stim_id = stimuli.query('first_30_chars == @first_30_chars').index
        assert len(stim_id) == 1
        stim_id = stim_id[0]

        n_words = len(stimuli.loc[stim_id, 'sequence'].split(' '))
        if (n_words != sum(sel)) and verbose:
            print('mistach of %i words in %s (stim %i)' % (n_words - sum(sel),
                                                           pos, stim_id))
            print('stim: %s' % stimuli.loc[stim_id, 'sequence'])
            print('log: %s' % ' '.join(row.loc[sel, 'word']))

        # Update
        log.loc[row.index, 'stim_id'] = stim_id
    return log

def read_log(log_fname, task='auto', verbose=False):
    log = _parse_log(log_fname)
    log = _clean_log(log)
    if task == 'auto':
        task = 'visual' if log_fname[-7:] == 'Vis.log' else 'auditory'
    if task == 'visual':
        # TODO: add sequence annotation for auditory
        log = _add_stim_id(log, verbose=verbose)
    return log

def get_log_times(log, events, sfreq):
    sel = np.sort(np.r_[
        np.where(events[:, 2] == 20)[0],  # fixation
        np.where(events[:, 2] == 10)[0]  # context
    ])
    common_megs = events[sel]
    common_logs = log.query('(new_context == True) or condition=="fix"')

    last_log = common_logs.time[0]
    last_meg = common_megs[0, 0]
    last_idx = 0
    assert len(common_megs) == len(common_logs)
    for common_meg, (idx, common_log) in zip(
            common_megs, common_logs.iterrows()):

        if common_meg[2] == 20:
            assert common_log.condition == 'fix'
        else:
            assert common_log.condition in ('sentence', 'word_list')

        log.loc[idx, 'meg_time'] = common_meg[0] / sfreq

        sel = slice(last_idx+1, idx)
        times = log.loc[sel, 'time'] - last_log + last_meg / sfreq
        assert np.all(np.isfinite(times))
        log.loc[sel, 'meg_time'] = times

        last_log = common_log.time
        last_meg = common_meg[0]
        last_idx = idx

        assert np.isfinite(last_log) * np.isfinite(last_meg)

    # last block
    sel = slice(last_idx+1, None)
    times = log.loc[sel, 'time'] - last_log + last_meg / sfreq
    log.loc[sel, 'meg_time'] = times
    log['meg_sample'] = np.array(log.meg_time.values * sfreq, int)
    return log

def match_list(A, B, on_replace='delete'):
    """Match two lists of different sizes and return corresponding indice
    Parameters
    ----------
    A: list | array, shape (n,)
        The values of the first list
    B: list | array: shape (m, )
        The values of the second list
    Returns
    -------
    A_idx : array
        The indices of the A list that match those of the B
    B_idx : array
        The indices of the B list that match those of the A
    """
    from Levenshtein import editops

    unique = np.unique(np.r_[A, B])
    label_encoder = dict((k, v) for v, k in enumerate(unique))

    def int_to_unicode(array):
        return ''.join([str(chr(label_encoder[ii])) for ii in array])

    changes = editops(int_to_unicode(A), int_to_unicode(B))
    B_sel = np.arange(len(B)).astype(float)
    A_sel = np.arange(len(A)).astype(float)
    for type, val_a, val_b in changes:
        if type == 'insert':
            B_sel[val_b] = np.nan
        elif type == 'delete':
            A_sel[val_a] = np.nan
        elif on_replace == 'delete':
            # print('delete replace')
            A_sel[val_a] = np.nan
            B_sel[val_b] = np.nan
        elif on_replace == 'keep':
            # print('keep replace')
            pass
        else:
            raise NotImplementedError
    B_sel = B_sel[np.where(~np.isnan(B_sel))]
    A_sel = A_sel[np.where(~np.isnan(A_sel))]
    assert len(B_sel) == len(A_sel)
    return A_sel.astype(int), B_sel.astype(int)

def add_part_of_speech(df):
    sentences = ' '.join(df['word'].values)
    while '  ' in sentences:
        sentences = sentences.replace('  ', ' ')

    nlp = spacy.load("nl_core_news_sm")
    doc = [i for i in nlp(sentences)]
    from_idx, to_idx = match_list([i.text for i in doc],
                                  df.word.values)

    part_of_speech = [doc[i].pos_ for i in from_idx]
    idx = df.index.values[to_idx]
    df.loc[idx, 'pos'] = part_of_speech
    df.loc[df.pos == 'X', 'pos'] = pd.np.nan
    #categories = ['ADJ', 'ADP', 'ADV', 'CONJ', 'DET',
                  # 'NOUN', 'NUM', 'PRON', 'PROPN', 'VERB']
    #df = df.join(pd.get_dummies(df.pos, columns=categories))
    df = df.join(pd.get_dummies(df.pos))
    if 'PROPN' not in df.keys():
        verb_idx = df.columns.get_loc("VERB")
        df.insert(verb_idx, 'PROPN', 0)

    return df

def add_word_frequency(df):
    freq = df.word.apply(lambda word: word_frequency(word, 'nl'))  # noqa
    df['word_freq'] = freq
    return df

def add_word_length(df):
    df['word_length'] = df.word.astype(str).apply(len)
    return df

def add_letter_count(df):
    word = df.word.str
    for letter in 'abcdefghijklmnopqrstuvwxyz':
        df[letter] = word.count(letter)
    return df

def read_mri_events(event_fname):
    """This is needs to be enriched depending on the analysis
    """
    # Read MRI events
    events = pd.read_csv(event_fname, sep='\t')

    # Add context: sentence or word list?
    contexts = dict(WOORDEN='word_list', ZINNEN='sentence')
    for key, value in contexts.items():
        sel = events.value.str.contains(key)
        events.loc[sel, 'context'] = value
        events.loc[sel, 'condition'] = value

    # Clean up MRI event mess
    sel = ~events.context.isna()
    start = 0
    context = 'init'
    for idx, row in events.loc[sel].iterrows():
        events.loc[start:idx, 'context'] = context
        start = idx
        context = row.context
    events.loc[start:, 'context'] = context

    # Add event condition: word, blank, inter stimulus interval etc
    conditions = (('50', 'pulse'), ('blank', 'blank'), ('ISI', 'isi'))
    for key, value in conditions:
        sel = events.value == key
        events.loc[sel, 'condition'] = value

    events.loc[events.value.str.contains('FIX '), 'condition'] = 'fix'

    # Extract words from file
    sel = events.condition.isna()
    words = events.loc[sel, 'value'].apply(lambda s: s.strip('0123456789 '))
    events.loc[sel, 'word'] = words

    # Remove empty words
    sel = (events.word.astype(str).apply(len) == 0) & (events.condition.isna())
    events.loc[sel, 'word'] = pd.np.nan
    events.loc[sel, 'condition'] = 'blank'
    events.loc[~events.word.isna(), 'condition'] = 'word'

    # --- Add word frequency
    sel = events.condition == 'word'

    def get_word_freq(word):
        return word_frequency(word, 'en', wordlist='best', minimum=0.0)

    events.loc[sel, 'word_freq'] = events.loc[sel, 'word'].apply(get_word_freq)

    # --- Add word length
    sel = events.condition == 'word'
    events['word_length'] = events.word.astype(str).apply(len)

    # --- TODO Add whatever features may be relevant here

    return events


from django.core.management.base import BaseCommand, CommandError
from django.conf import settings
from django.core.mail import EmailMessage

import logging
logger = logging.getLogger(__name__)

from notifications.models import Subscription

FROM_EMAIL = getattr(settings, 'NOTIFICATION_FROM_EMAIL', 'notifications@wheresyourtrash.com')

class Command(BaseCommand):
    help = 'Run through subscriptions sending out notifications'

    def handle(self, *args, **options):
        # Gather all subscriptions that need to go out
        #subs = Subscription.objects.live()
        for sub in Subscription.objects.all():
            if sub.day_of_pickup:
                logger.info('Sending notification to {0}'.format(sub.user))
                pickup_type = sub.district.district_type.lower() # + '_notification'
                if sub.subscription_type == 'SMS':
                    recipient = sub.clean_phone_number + '@' + sub.service_provider.email_root
                else:
                    recipient = sub.user.email

                if sub.district.district_type == 'TRASH':
                    pickup_str = 'trash'
                else:
                    pickup_str = 'recycling'
                txt_content = 'Is your {0} outside and ready? The folks from {1} are going to be picking it up later today!!'.format(pickup_str, sub.district.municipality.name)
                email = EmailMessage("Where's Your Trash?", txt_content, FROM_EMAIL, [recipient])
                email.send()

                '''
                send_templated_mail(template_name=template_name,
                                    from_email=FROM_EMAIL,
                                    recipient_list=[recipient],
                                    context={'subscription':sub})
                '''


from django.contrib import admin

from django.contrib.admin import AdminSite
from django.http import HttpResponse, HttpResponseRedirect

from .models import Category,Item,SalesMeasure,Presentation,OrderItems,Order,Customer,Payment,Shop
from django.shortcuts import render,redirect
from django.contrib.admin import *

from GorceryDelivery.forms import OrderDisplayForm
from GorceryDelivery.modules.PDFGenerator import *
from GorceryDelivery.modules.AggregateOrderPDF import *
from datetime import datetime

class MyAdminSite(AdminSite):

    def __init__(self, *args, **kwargs):
        """
        allows complete inheritance from the default admin: default registration of USER and GROUP in admin
        :param args:
        :param kwargs:
        """
        super(MyAdminSite, self).__init__(*args, **kwargs)
        self._registry.update(site._registry)  # PART 2

    def get_urls(self):
        from django.urls import path,re_path
        urls = super().get_urls()
        urls += [

            path('customer_orders/', self.admin_view(self.get_daily_orders)),
            path('customer_orders/<str:delivery_date>/', self.admin_view(self.display_daily_orders)),
            re_path(r'^download_single_order/(\d+)/$', self.admin_view(self.download_single_order)),
            re_path(r'^customer_orders/details/(\d+)/$', self.admin_view(self.display_order_details)),
            path('customer_orders/items_by_day/<str:delivery_date>/', self.admin_view(self.aggregate_daily_orders))
            #path('customer_orders/print_all/<str:delivery_date>/', self.admin_view(self.print_orders_by_day))
        ]
        return urls

    def get_daily_orders(self, request):
        """
        View function: allows displays of all order given a specific delivery date
        :param request: Http request
        :return: Http redirect to display page; if method is not 'POST', renders delivery date form
        """

        if request.method == 'POST':
            form = OrderDisplayForm(request.POST)
            if form.is_valid():
                cd = form.cleaned_data
                next_url = '/admin/customer_orders/' + cd.get('date') + '/'
                #assert False
                return HttpResponseRedirect(next_url)# redirect to url of 'display daily orders'

        else:
            form = OrderDisplayForm()

        return  render(request, 'admin/get_orders.html', {'form': form})

    def display_daily_orders(self, request, delivery_date):
        """
        View function: actually dispays the orders
        :param request: Http request
        :param delivery_date: delivery date to query
        :return: renders a tabled list of all the orders
        """

        date_error = False
        # convert the date first
        try:#since date input is exposed using get: anticipating incorrect input
            objDate = datetime.strptime(delivery_date, '%d-%m-%Y')
            delivery_date = objDate.strftime('%Y-%m-%d')
        except:
            delivery_date = None
            date_error = True

        date_orders = Order.objects.filter(delivery_date=delivery_date)

        return render(request, 'admin/show_orders.html', {'orders': date_orders,
                                                          'date': delivery_date,
                                                          'date_error': date_error})

    def download_single_order(self, request, order_Id):
        """
        View function: downloads the user order form the 'display daily orders' list as pdf
        :param request: Http request
        :param order_Id: pk of the order in db
        :return: pdf file as Http response
        """

        bought_items = None
        order_details = None

        order_Id = int(order_Id)

        try:
            order_details = Order.objects.get(id=order_Id)
        except:
            order_details = None

        bought_items = OrderItems.objects.filter(order=order_details)

        if not bought_items == None:
            the_report = PDFGenerator(request,order_details, bought_items)
            return the_report.download_user_order()

    def display_order_details(self, request, order_id):
        """
        View function: displays the details of a selected custmer order
        :param request: Http request
        :param order_id: pk of the selected order
        :return: renders the details
        """

        order_details = None
        bought_items = None
        order_id = int(order_id)

        try:
            order_details = Order.objects.get(id=order_id)
        except:
            order_details = None

        bought_items = OrderItems.objects.filter(order=order_details)

        return render(request,'admin/display_order_details.html', { 'order': order_details,
                                                                    'bought_items': bought_items})

    def aggregate_daily_orders(self, request, delivery_date):
        """
        View Function: downloads aggregated orders for a given date showi ng an item and how many times it was ordered.
        :param request: Http request
        :param delivery_date: delivery date
        :return: pdf file as http response
        """

        date_error = False
        try:
            objDate = datetime.strptime(delivery_date, '%Y-%m-%d')
            delivery_date = objDate.strftime('%Y-%m-%d')
        except:
            date_error = True

        order_items_list = []

        if not date_error:
            order_list = Order.objects.filter(delivery_date = delivery_date)
            for order in order_list:
                item_list = OrderItems.objects.filter(order= order)
                order_items_list.append(item_list)

            daily_order = AggregateOrderPDF(order_items_list, delivery_date)
            daily_order.compute_aggregate_list()
            return  daily_order.download_aggregate_order()

        #assert False

admin_site = MyAdminSite()

@admin.register(Item, site= admin_site)
class AdminItem(admin.ModelAdmin):
    list_display = ('name', 'description', 'price', 'discounted_price', 'availability')
    list_filter = ('category', 'availability')
    search_fields = ('name',)
    list_per_page = 10

@admin.register(Category, site =admin_site)
class AdminCategory(admin.ModelAdmin):
    list_display = ('name', 'description')
    list_per_page = 10
    search_fields = ('name',)

@admin.register(Customer, site =admin_site)
class AdminCustomer(admin.ModelAdmin):
    list_display = ('user_profile','gender', 'phone_no', 'preffered_delivery_addr')
    list_per_page = 10

@admin.register(Order, site=admin_site)
class AdminOrder(admin.ModelAdmin):
    list_display =('buyer_firstname', 'buyer_lastname', 'buyer_gender','buyer_email', 'buyer_phone_no',
                   'order_date', 'paid_for','delivery_date','order_total','delivery_addr', 'customer' )
    list_per_page = 10
    list_filter = ('buyer_gender', 'paid_for', 'order_date', 'delivery_date')
    search_fields = ( 'buyer_firstname', 'buyer_lastname','buyer_phone_no','buyer_email')

@admin.register(SalesMeasure, site =admin_site)
class AdminSalesMeasure(admin.ModelAdmin):
    list_display = ('name','description', 'item', 'price')
    list_per_page = 10
    search_fields = ('item__name','name')
    list_filter = ('item',)

@admin.register(Presentation, site = admin_site)
class AdminPresentation(admin.ModelAdmin):
    list_per_page = 10
    list_display = ('name', 'description', 'item')

@admin.register(Payment, site = admin_site)
class AdminPayment(admin.ModelAdmin):
    list_display = ('order', 'payment_amount',  'timestamp', 'transaction_reference', 'status','description' )
    list_per_page = 10

    def payment_amount(self, obj):
        return obj.order.order_total

@admin.register(OrderItems, site=admin_site)
class AdminOrderItem(admin.ModelAdmin):
    list_display = ('order_item_desc', 'order_item_price', 'quantity', 'price_by_qty')
    list_per_page = 10
    list_filter = ('item', 'order')
    search_fields = ('item__name', 'order__buyer_firstname')

@admin.register(Shop, site=admin_site)
class Shop(admin.ModelAdmin):
    list_display = ('shop_name','site_domain', 'address', 'email', 'phone_number', 'phone_number_2', 'is_active')

{'path': 'django.urls.path', 're_path': 'django.urls.re_path'}


from selenium import webdriver
from selenium.webdriver.chrome.options import Options

def main():
    chrome_options = Options()
    chrome_options.add_argument("--disable-extensions")
    chrome_options.add_argument("--disable-gpu")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--headless")
    chrome_options.headless = True  # also works
    driver = webdriver.Chrome(executable_path='/Users/jorgecorradi/Downloads/chromedriver', options=chrome_options)
    driver.get('https://www.infomoney.com.br/cotacoes/empresas-b3/')

    symbols = []
    for el in driver.find_elements_by_css_selector('.list-companies a'):
        symbols.append(el.text)

    with open('/Users/jorgecorradi/Desktop/stocks.csv', 'w') as csv_file:
        csv_file.write("symbol;companyName;stockType;sector;category;currency" + "\n")
        for symbol in symbols:
            driver.get('https://www.fundamentus.com.br/detalhes.php?papel=' + symbol)
            print("Scraping " + symbol)

            try:
                data = {}
                data['"symbol"'] = symbol

                type = driver.find_element_by_css_selector(
                    'body > div.center > div.conteudo.clearfix > table:nth-child(2) > tbody > tr:nth-child(1) > td.label.w15 > span.txt').text

                if type == 'Papel':
                    data['companyName'] = driver.find_element_by_css_selector(
                        'body > div.center > div.conteudo.clearfix > table:nth-child(2) > tbody > tr:nth-child(3) > td:nth-child(2) > span').text

                    data['stockType'] = driver.find_element_by_css_selector(
                        'body > div.center > div.conteudo.clearfix > table:nth-child(2) > tbody > tr:nth-child(2) > td:nth-child(2) > span').text
                    data['sector'] = driver.find_element_by_css_selector(
                        'body > div.center > div.conteudo.clearfix > table:nth-child(2) > tbody > tr:nth-child(4) > td:nth-child(2) > span > a').text
                    data['category'] = 'STOCK'
                else:
                    data['companyName'] = driver.find_element_by_css_selector(
                        'body > div.center > div.conteudo.clearfix > table:nth-child(3) > tbody > tr:nth-child(2) > td:nth-child(2) > span').text

                    data['stockType'] = 'FII'
                    data['sector'] = driver.find_element_by_css_selector(
                        'body > div.center > div.conteudo.clearfix > table:nth-child(3) > tbody > tr:nth-child(4) > td:nth-child(2) > span > a').text
                    data['category'] = 'REF'

                data['currency'] = 'BRL'

                csv_file.write(";".join(data.values()) + "\n")
                csv_file.flush()
            except:
                print("couldn't get info for " + symbol)

if __name__ == '__main__':
    main()
    print("Done Scraping")


from flask_restful import Resource, reqparse, abort, fields, marshal_with
from datetime import datetime
from . import citizen_api
from ..database_models import Citizen
from .. import data_handler, sms_handler
from app import db

# Parser to check that required arguments are sent to add a citizen to the database
citizen_put_args = reqparse.RequestParser()
citizen_put_args.add_argument("id_citizen", type=str, required=True,
                              help="The ID of the citizen is required. The argument should be a String.")
citizen_put_args.add_argument("email", type=str, required=True,
                              help="The email of the citizen is required. The argument should be a String.")
citizen_put_args.add_argument("name", type=str, required=True,
                              help="The name of the citizen is required. The argument should be a String.")
citizen_put_args.add_argument("surname", type=str, required=True,
                              help="The surname of the citizen is required. The argument should be a String.")
citizen_put_args.add_argument("date_of_birth", type=str, required=True,
                              help="The date of birth of the citizen is required in the format YYYY-MM-DD. "
                                   "The argument should be a String.")
citizen_put_args.add_argument("mobile_num", type=str, required=True,
                              help="The mobile number of the citizen is required. The argument should be a String.")
citizen_put_args.add_argument("medical_aid", type=str,
                              help="The medical aid number of the citizen should be a String.")
citizen_put_args.add_argument("citizen_address", type=str,
                              help="The citizen address should be a String.")

# Parser to check that required arguments are sent to get a citizen from the database
citizen_get_args = reqparse.RequestParser()
citizen_get_args.add_argument("id_citizen", type=str, required=True,
                              help="The ID of the citizen is required. The argument should be a String.")

# Parser to check that required arguments are sent to update a citizen in the database
citizen_patch_args = reqparse.RequestParser()
citizen_patch_args.add_argument("id_citizen", type=str, required=True,
                                help="The ID of the citizen is required. The argument should be a String.")
citizen_patch_args.add_argument("email", type=str,
                                help="The citizen email address should be a String.")
citizen_patch_args.add_argument("name", type=str,
                                help="The citizen name should be a String.")
citizen_patch_args.add_argument("surname", type=str,
                                help="The citizen surname should be a String.")
citizen_patch_args.add_argument("date_of_birth", type=str,
                                help="The date of birth of the citizen is should be in the format YYYY-MM-DD. "
                                     "The argument should be a String.")
citizen_patch_args.add_argument("mobile_num", type=str,
                                help="The mobile number of the citizen should be a String")
citizen_patch_args.add_argument("medical_aid", type=str,
                                help="The medical aid number of the citizen should be a String.")
citizen_patch_args.add_argument("citizen_address", type=str,
                                help="Address is required in string format")

# Parser to check that required arguments are sent to delete a citizen from the database
citizen_del_args = reqparse.RequestParser()
citizen_del_args.add_argument("id_citizen", type=str, required=True,
                              help="The ID of the citizen is required. The argument should be a String.")

# Fields to marshal the responses
resource_fields = {
    'id_citizen': fields.String,
    'email': fields.String,
    'name': fields.String,
    'surname': fields.String,
    'date_of_birth': fields.String,
    'mobile_num': fields.String,
    'medical_aid': fields.String,
    'citizen_address': fields.String
}

# Class to handle methods related to citizens
class CitizenResource(Resource):
    @marshal_with(resource_fields)
    def get(self):
        args = citizen_get_args.parse_args()
        data_handler.clean_data(args)
        return get_citizen(args), 200

    def put(self):
        args = citizen_put_args.parse_args()
        data_handler.clean_data(args)
        add_citizen(args)
        return {"message": "Added to database"}, 201

    def patch(self):
        args = citizen_patch_args.parse_args()
        data_handler.clean_data(args)
        update_citizen(args)
        return {"message": "Updated the database"}, 200

    def delete(self):
        args = citizen_del_args.parse_args()
        data_handler.clean_data(args)
        delete_citizen(args)
        return {"message": "Deleted from database"}, 204

# Add resource to the API
citizen_api.add_resource(CitizenResource, "")

# Get a citizen by their ID and return the full database entry
def get_citizen(args):
    result = Citizen.query.filter_by(id_citizen=args["id_citizen"]).first()
    if result:
        return result
    else:
        abort(404, message="A citizen with this ID does not exist")

# Add a citizen to the database
def add_citizen(args):
    result = Citizen.query.filter_by(id_citizen=args["id_citizen"]).first()
    if result:
        abort(409, message="A citizen with this ID already exists")
    else:
        try:
            date = datetime.strptime(args["date_of_birth"], '%Y-%m-%d')
            new_citizen = Citizen(id_citizen=args["id_citizen"], email=args["email"], name=args["name"],
                                  surname=args["surname"], date_of_birth=date, mobile_num=args["mobile_num"],
                                  medical_aid=args["medical_aid"],
                                  citizen_address=args["citizen_address"])
            db.session.add(new_citizen)
            db.session.commit()
            sms_handler.send_mock_sms("You have been added", args["mobile_num"])
        except ValueError:
            abort(400, message="The date should be in format YYYY-MM-DD")

# Update the citizen in the database
def update_citizen(args):
    result = Citizen.query.filter_by(id_citizen=args["id_citizen"]).first()
    if not result:
        abort(404, message="Citizen with this ID does not exist, cannot update")
    else:
        if args["email"]:
            result.email = args["email"]
        if args["name"]:
            result.name = args["name"]
        if args["surname"]:
            result.surname = args["surname"]
        if args["date_of_birth"]:
            try:
                date = datetime.strptime(args["date_of_birth"], '%Y-%m-%d')
                result.date_of_birth = date
            except ValueError:
                abort(400, message="The date should be in format YYYY-MM-DD")
        if args["mobile_num"]:
            result.mobile_num = args["mobile_num"]
        if args["medical_aid"]:
            result.medical_aid = args["medical_aid"]
        if args["citizen_address"]:
            result.citizen_address = args["citizen_address"]
        db.session.commit()

# Delete a citizen from the database
def delete_citizen(args):
    result = Citizen.query.filter_by(id_citizen=args["id_citizen"]).first()
    if not result:
        abort(404, message="Citizen with this ID does not exist, cannot delete")
    else:
        db.session.delete(result)
        db.session.commit()


# coding=UTF-8
#读取config.ini配置信息中的数据

import os
import codecs
import configparser
#将读取配置信息ini的过程封装成一个类
proDir = os.path.split(os.path.realpath(__file__))[0]    #该文件的地址，这里相当于这里的接口测试文档框架路径,绝对路径
configPath = os.path.join(proDir, "config.ini")         #存放在根目录下文件名是config.ini的文件，配置文件的完整地址

class ReadConfig():
    def __init__(self):
        fd = open(configPath ,encoding="utf-8")          #打开配置文件
        data = fd.read()               #读取配置文件

        #  remove BOM                  #判断文件中是否有数据
        if data[:3] == codecs.BOM_UTF8:
            data = data[3:]
            file = codecs.open(configPath, "w")
            file.write(data)
            file.close()
        fd.close()

        self.cf = configparser.ConfigParser()   #调用读取配置模块的类
        self.cf.read(configPath,encoding="utf-8")                #读取文件

#获取emall、http、database分组下指定的name值
    def get_email(self, name):
        value = self.cf.get("EMAIL", name)
        return value

    def get_http(self, name):
        value = self.cf.get("HTTP", name)
        return value

    def get_headers(self, name):
        value = self.cf.get("HEADERS", name)
        return value

    def set_headers(self, name, value):
        self.cf.set("HEADERS", name, value)
        with open(configPath, 'w+') as f:
            self.cf.write(f)

    def get_url(self, name):
        value = self.cf.get("URL", name)
        return value

    def get_db(self, name):
        value = self.cf.get("DATABASE", name)
        return value

#修改http分组下指定name的值value   现阶段还未使用
    def set_http(self,name,value):
        cfg = self.cf.set("HTTP",name,value)
        fp = open(r'config.ini','w')
        cfg.write(fp)


import argparse
import csv
import math
import os
import os.path as osp

import joblib
import numpy as np
import src.ast_toolbox.mcts.BoundedPriorityQueues as BPQ
# from example_save_trials import *
import tensorflow as tf
from CartPole.cartpole_simulator import CartpoleSimulator
# from garage.tf.algos.trpo import TRPO
from garage.baselines.linear_feature_baseline import LinearFeatureBaseline
from garage.misc import logger
from garage.tf.policies.gaussian_mlp_policy import GaussianMLPPolicy
from src.ast_toolbox import TRPO
from src.ast_toolbox import ASTEnv
from src.ast_toolbox import TfEnv
from src.ast_toolbox.rewards import ASTRewardS

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # just use CPU

# Logger Params
parser = argparse.ArgumentParser()
parser.add_argument('--exp_name', type=str, default='cartpole')
parser.add_argument('--nd', type=int, default=1)
parser.add_argument('--sut_itr', type=int, default=5)
parser.add_argument('--n_trial', type=int, default=10)
parser.add_argument('--trial_start', type=int, default=0)
parser.add_argument('--n_itr', type=int, default=1000)
parser.add_argument('--step_size', type=float, default=5.0)
parser.add_argument('--batch_size', type=int, default=2000)  # 50000
parser.add_argument('--snapshot_mode', type=str, default="none")
parser.add_argument('--snapshot_gap', type=int, default=500)
parser.add_argument('--log_dir', type=str, default='./Data/AST/TRPO')
parser.add_argument('--args_data', type=str, default=None)
args = parser.parse_args()
args.log_dir += ('B' + str(args.batch_size) + 'Step' + str(args.step_size))

top_k = 10
max_path_length = 100
open_loop = False

tf.set_random_seed(0)
sess = tf.Session()
sess.__enter__()

# Instantiate the env
data = joblib.load("../CartPole/ControlPolicy/itr_" + str(args.sut_itr) + ".pkl")
sut = data['policy']
reward_function = ASTRewardS()

simulator = CartpoleSimulator(sut=sut, max_path_length=100, use_seed=False, nd=args.nd)
env = TfEnv(ASTEnv(open_loop=open_loop,
                   simulator=simulator,
                   fixed_init_state=True,
                   s_0=[0.0, 0.0, 0.0 * math.pi / 180, 0.0],
                   reward_function=reward_function,
                   ))

# Create policy
policy = GaussianMLPPolicy(
    name='ast_agent',
    env_spec=env.spec,
    hidden_sizes=(128, 64, 32)
)

with open(osp.join(args.log_dir, 'total_result.csv'), mode='w') as csv_file:
    fieldnames = ['step_count']
    for i in range(top_k):
        fieldnames.append('reward ' + str(i))
    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
    writer.writeheader()

    for trial in range(args.trial_start, args.trial_start + args.n_trial):
        # Create the logger
        log_dir = args.log_dir + '/' + str(trial)

        tabular_log_file = osp.join(log_dir, 'process.csv')
        text_log_file = osp.join(log_dir, 'text.txt')
        params_log_file = osp.join(log_dir, 'args.txt')

        logger.set_snapshot_dir(log_dir)
        logger.set_snapshot_mode(args.snapshot_mode)
        logger.set_snapshot_gap(args.snapshot_gap)
        logger.log_parameters_lite(params_log_file, args)
        if trial > args.trial_start:
            old_log_dir = args.log_dir + '/' + str(trial - 1)
            logger.pop_prefix()
            logger.remove_text_output(osp.join(old_log_dir, 'text.txt'))
            logger.remove_tabular_output(osp.join(old_log_dir, 'process.csv'))
        logger.add_text_output(text_log_file)
        logger.add_tabular_output(tabular_log_file)
        logger.push_prefix("[" + args.exp_name + '_trial ' + str(trial) + "]")

        np.random.seed(trial)

        params = policy.get_params()
        sess.run(tf.variables_initializer(params))
        baseline = LinearFeatureBaseline(env_spec=env.spec)
        # optimizer = ConjugateGradientOptimizer(hvp_approach=FiniteDifferenceHvp(base_eps=1e-5))

        top_paths = BPQ.BoundedPriorityQueue(top_k)
        algo = TRPO(
            env=env,
            policy=policy,
            baseline=baseline,
            batch_size=args.batch_size,
            step_size=args.step_size,
            n_itr=args.n_itr,
            store_paths=True,
            # optimizer= optimizer,
            max_path_length=max_path_length,
            top_paths=top_paths,
            plot=False,
        )

        algo.train(sess=sess, init_var=False)

        row_content = dict()
        row_content['step_count'] = args.n_itr * args.batch_size
        i = 0
        for (r, action_seq) in algo.top_paths:
            row_content['reward ' + str(i)] = r
            i += 1
        writer.writerow(row_content)


import wx

#----------------------------------------------------------------------

class TestPanel(wx.Panel):
    def __init__(self, parent, log):
        self.log = log
        wx.Panel.__init__(self, parent, -1)

        dpc = wx.DatePickerCtrl(self, size=(120,-1),
                                style=wx.DP_DROPDOWN | wx.DP_SHOWCENTURY)

        sizer = wx.BoxSizer(wx.VERTICAL)
        sizer.Add(dpc, 0, wx.ALL, 50)

        self.SetSizer(sizer)

#----------------------------------------------------------------------

def runTest(frame, nb, log):
    win = TestPanel(nb, log)
    return win

#----------------------------------------------------------------------

overview = """

wx.DatePickerCtrl

This control allows the user to select a date. Unlike
wx.calendar.CalendarCtrl, which is a relatively big control,
wx.DatePickerCtrl is implemented as a small window showing the
currently selected date. The control can be edited using the keyboard,
and can also display a popup window for more user-friendly date
selection, depending on the styles used and the platform.

"""

if __name__ == '__main__':
    import sys,os
    import run
    run.main(['', os.path.basename(sys.argv[0])] + sys.argv[1:])


import cv2
import numpy as np

filename = 'cup.jpg'

print('顯示圖片')
image = cv2.imread(filename, -1)

shape = image.shape
h = shape[0]    #高
w = shape[1]    #寬
h, w, d = image.shape   #d為dimension d=3 全彩 d=1 灰階
print("寬 = ", w, ", 高 = ", h, ", D = ", d)

image = cv2.resize(image, (403, 302))
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray = cv2.GaussianBlur(gray, (5, 5), 0)
circles = cv2.HoughCircles(
    gray,
    cv2.HOUGH_GRADIENT, # 偵測方法目前只支援這個參數
    1,      # 1代表偵測圖與輸入圖大小一致，填1即可
    20,     # 各圓心間的最小距離，設太小容易誤判，太大會將數個圓當成一個
    None,   # 固定填 None
    10,     # canny演算法的高閾值，此值一半為低閾值
    75,     # 超過此閾值才會被當作圓
    3,      # 最小圓半徑
    75      # 最大圓半徑
)

circles = circles.astype(int)
if len(circles) > 0:
    out = image.copy()
    for x, y, r in circles[0]:
        # 畫圓
        cv2.circle(out, (x, y), r, (0, 0, 255), 3)
        # 畫圓心
        cv2.circle(out, (x, y), 2, (0, 255, 0), 3)
    image = cv2.hconcat([image, out])

cv2.namedWindow('image', cv2.WINDOW_NORMAL)
cv2.imshow('image', image) #顯示圖片

print('在此等待任意鍵繼續, 繼續後刪除本視窗')
cv2.waitKey()
cv2.destroyAllWindows()


"""Ciphers module.
Provide Asymmetric and Symmetric cipher functions.\
"""
from enum import IntEnum

from Crypto import Random
from Crypto.Cipher import PKCS1_v1_5, AES
from Crypto.PublicKey import RSA
from Crypto.Random import random as rand
from Crypto.Util import Counter

class AESBlockCipher:

    class Size(IntEnum):
        bit_128 = 1
        bit_192 = 2
        bit_256 = 3

    supported_mode = [AES.MODE_CBC, AES.MODE_CTR]

    def __init__(self, keysize=Size.bit_256, mode=AES.MODE_CBC, secret=None, salt=None):
        """Generate AES key according to the keysize and secret.
        keysize -- int, AES keysize in byte (16, 24, 32).
        secret -- bytes, only used for Password-Based AES cipher.
        """

        if mode not in self.supported_mode:
            raise None  # TODO raise error

        self.aes = None
        self.MODE = mode

        if keysize not in list(map(int, self.Size)):
            raise Exception("Invalid key size.")

        keysize = int((128 + 64 * (keysize - 1)) / 8)

        if not secret:
            """Generate randome key"""
            self.key = Random.get_random_bytes(keysize)
        else:  # Password-Based
            """Construct AES key from user secret with argon2 hash function"""
            import argon2

            self.salt = Random.get_random_bytes(16) if salt is None else salt
            time_cost = 2
            memory_cost = 512
            parallelism = 2
            hash_len = keysize

            self.key = argon2.low_level.hash_secret_raw(secret, self.salt, time_cost,
                                                        memory_cost, parallelism,
                                                        hash_len, argon2.low_level.Type.I)

    def _setup_cipher(self, iv=None):
        # 16-byte, 128-bit initial vector/value for both CBC and CTR mode
        iv = Random.get_random_bytes(AES.block_size) if not iv else iv
        if self.MODE == AES.MODE_CBC:
            return iv, AES.new(self.key, AES.MODE_CBC, iv)
        if self.MODE == AES.MODE_CTR:
            ctr = Counter.new(128, initial_value=int.from_bytes(iv, byteorder="big"))
            return iv, AES.new(self.key, AES.MODE_CTR, counter=ctr)

    @staticmethod
    def pad(b):
        return b + \
               ((AES.block_size - len(b) % AES.block_size) *
                chr(AES.block_size - len(b) % AES.block_size)).encode()

    @staticmethod
    def unpad(b):
        return b[:-b[-1]]

    # Encrypt a chuck of bytes
    def encrypt(self, b, last=False):
        iv = bytes()
        if not self.aes:  # first time, setup iv
            iv, self.aes = self._setup_cipher()

        if not last:
            return iv + self.aes.encrypt(b)
        else:
            # if last chunk, pad the chunk
            ret = self.aes.encrypt(self.pad(b))
            self.aes = None
            return iv + ret  # return iv also since it could be the only chunk

    def decrypt(self, b, last=False):
        if not self.aes:  # first time, retrieve iv
            iv = b[:AES.block_size]  # both modes have the same iv size
            b = b[len(iv):]
            _, self.aes = self._setup_cipher(iv)

        if not last:
            return self.aes.decrypt(b)
        else:
            # if last chunk, unpad the chunk
            dec = self.aes.decrypt(b)
            return self.unpad(dec)

class RSACipher:
    rsa = None
    cipher = None

    class Size(IntEnum):
        bit_1024 = 1
        bit_2048 = 2
        bit_4096 = 3
        bit_8192 = 4

    def import_key(self, key, passphrase=None):
        # if key is bytes, load it directly
        if isinstance(key, bytes):
            self.rsa = RSA.importKey(key, passphrase)
        else:  # else read it as file
            with open(key, "rb") as f:
                self.rsa = RSA.importKey(f.read(), passphrase)

        self.cipher = PKCS1_v1_5.new(self.rsa)

    def export_key(self, keyfile=None, format="PEM", passphrase=None, public_only=False):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        # if has keyfile, write to key file
        if keyfile:
            with open(keyfile, "wb") as f:
                if public_only:
                    pubkey = self.rsa.publickey()
                    f.write(pubkey.exportKey(format=format))
                else:
                    f.write(self.rsa.exportKey(format=format, passphrase=passphrase))
        else:  # else return bytes
            if public_only:
                pubkey = self.rsa.publickey()
                return pubkey.exportKey(format=format)
            else:
                return self.rsa.exportKey(format=format, passphrase=passphrase)

    def generate(self, keysize=Size.bit_4096, progress_func=None):
        self.progress_func = progress_func

        if keysize not in list(map(int, self.Size)):
            raise Exception("Invalid key size.")

        keysize = {1: 1024, 2: 2048, 3: 4096, 4: 8192}[keysize]
        self.rsa = RSA.generate(keysize, progress_func=self.gen_prog)

        self.cipher = PKCS1_v1_5.new(self.rsa)

        self.gen_prog("Finished, key size: " + str(RSACipher.Size(keysize)))

    # progress function for generating key
    def gen_prog(self, s):
        d = {"p,q": "Generating p and q primes...",
             "u": "Generating public exponent...",
             "d": "Generating private exponent..."}
        if self.progress_func:
            if s in d.keys():
                self.progress_func(d[s])
            else:
                self.progress_func(s)

    # encrypt with rsa key, data size should not exceed n - 11
    def encrypt(self, b):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        return self.cipher.encrypt(b)

    def decrypt(self, b):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        return self.cipher.decrypt(b, Random.get_random_bytes(rand.randint(0, self.rsa.size())))


from base.socialRecommender import SocialRecommender
from scipy.sparse import *
import numpy as np
from numpy import linalg as LA
from joblib import Parallel, delayed
from math import sqrt
EPS = 1e-8
# this model refers to the following paper:
# #########----  Collaborative Filtering with Social Exposure: A Modular Approach to Social Recommendation   ----#############
# SEREC_boost
class SERec(SocialRecommender):
    def __init__(self,conf,trainingSet=None,testSet=None,relation=None,fold='[1]'):
        super(SERec, self).__init__(conf,trainingSet,testSet,relation,fold)

    def initModel(self):
        super(SERec, self).initModel()
        self.lam_theta = 1e-5
        self.lam_beta = 1e-5
        self.lam_y = 0.01
        self.init_mu = 0.01
        self.a = 1.0
        self.b = 99.0
        self.s= 2.2
        self.init_std = 0.5
        self.theta = self.init_std * \
            np.random.randn(self.num_users, self.emb_size).astype(np.float32)
        self.beta = self.init_std * \
            np.random.randn(self.num_items, self.emb_size).astype(np.float32)
        self.mu = self.init_mu * np.ones((self.num_users,self.num_items), dtype=np.float32)
        self.n_jobs=4
        self.batch_size=1000
        row,col,val = [],[],[]
        for user in self.data.trainSet_u:
            for item in self.data.trainSet_u[user]:
                u = self.data.user[user]
                i = self.data.item[item]
                row.append(u)
                col.append(i)
                val.append(1)

        self.X = csr_matrix((np.array(val),(np.array(row),np.array(col))),(self.num_users,self.num_items))
        row,col,val = [],[],[]
        for user in self.social.followees:
            for f in self.social.followees[user]:
                u = self.data.user[user]
                i = self.data.user[f]
                row.append(u)
                col.append(i)
                val.append(1)
        self.T = csr_matrix((np.array(val), (np.array(row), np.array(col))), (self.num_users, self.num_users))

    def trainModel(self):
        print('training...')
        self._update(self.X)

    def _update(self, X):
        '''Model training and evaluation on validation set'''
        n_users = X.shape[0]
        XT = X.T.tocsr()  # pre-compute this
        self.vad_ndcg = -np.inf
        for i in range(self.maxEpoch):

            print('epoch #%d' % i)
            self._update_factors(X, XT)
            print(self.mu)
            self._update_expo(X, n_users)

    def _update_factors(self, X, XT):
        '''Update user and item collaborative factors with ALS'''
        self.theta = recompute_factors(self.beta, self.theta, X,
                                       self.lam_theta / self.lam_y,
                                       self.lam_y,
                                       self.mu,
                                       self.n_jobs,
                                       batch_size=self.batch_size)

        self.beta = recompute_factors(self.theta, self.beta, XT,
                                      self.lam_beta / self.lam_y,
                                      self.lam_y,
                                      self.mu,
                                      self.n_jobs,
                                      batch_size=self.batch_size)
    def _update_expo(self, X, n_users):
        '''Update exposure prior'''
        print('\tUpdating exposure prior...')
        start_idx = list(range(0, n_users, self.batch_size))
        end_idx = start_idx[1:] + [n_users]
        A_sum = np.zeros(self.num_items)
        for lo, hi in zip(start_idx, end_idx):
            A_sum += a_row_batch(X[lo:hi], self.theta[lo:hi], self.beta,
                                 self.lam_y, self.mu[lo:hi]).sum(axis=0)
        A_sum=np.tile(A_sum,[self.num_users,1])
        S_sum = self.T.dot(A_sum)
        self.mu = (self.a + A_sum +(self.s-1)*S_sum- 1) / (self.a + self.b + (self.s-1)*S_sum+n_users - 2)

    def predictForRanking(self,u):
        'invoked to rank all the items for the user'
        if self.data.containsUser(u):
            u = self.data.getUserId(u)
            return self.beta.dot(self.theta[u])
        else:
            return [self.data.globalMean] * self.num_items

def get_row(Y, i):
    '''Given a scipy.sparse.csr_matrix Y, get the values and indices of the
    non-zero values in i_th row'''
    lo, hi = Y.indptr[i], Y.indptr[i + 1]
    return Y.data[lo:hi], Y.indices[lo:hi]

def a_row_batch(Y_batch, theta_batch, beta, lam_y, mu):
    '''Compute the posterior of exposure latent variables A by batch'''
    pEX = sqrt(lam_y / 2 / np.pi) * \
          np.exp(-lam_y * theta_batch.dot(beta.T) ** 2 / 2)

    A = (pEX + EPS) / (pEX + EPS + (1 - mu) / mu)
    A[Y_batch.nonzero()] = 1.
    return A

def _solve(k, A_k, X, Y, f, lam, lam_y, mu):
    '''Update one single factor'''
    s_u, i_u = get_row(Y, k)
    a = np.dot(s_u * A_k[i_u], X[i_u])
    B = X.T.dot(A_k[:, np.newaxis] * X) + lam * np.eye(f)
    return LA.solve(B, a)

def _solve_batch(lo, hi, X, X_old_batch, Y, m, f, lam, lam_y, mu):
    '''Update factors by batch, will eventually call _solve() on each factor to
    keep the parallel process busy'''
    assert X_old_batch.shape[0] == hi - lo

    if mu.shape[1] == X.shape[0]:  # update users
        A_batch = a_row_batch(Y[lo:hi], X_old_batch, X, lam_y, mu[lo:hi])
    else:  # update items
        A_batch = a_row_batch(Y[lo:hi], X_old_batch, X, lam_y, mu.T[lo:hi])

    X_batch = np.empty_like(X_old_batch, dtype=X_old_batch.dtype)
    for ib, k in enumerate(range(lo, hi)):
        X_batch[ib] = _solve(k, A_batch[ib], X, Y, f, lam, lam_y, mu)
    return X_batch

def recompute_factors(X, X_old, Y, lam, lam_y, mu, n_jobs, batch_size=1000):
    '''Regress X to Y with exposure matrix (computed on-the-fly with X_old) and
    ridge term lam by embarrassingly parallelization. All the comments below
    are in the view of computing user factors'''
    m, n = Y.shape  # m = number of users, n = number of items
    assert X.shape[0] == n
    assert X_old.shape[0] == m
    f = X.shape[1]  # f = number of factors

    start_idx = list(range(0, m, batch_size))
    end_idx = start_idx[1:] + [m]
    res = Parallel(n_jobs=n_jobs)(delayed(_solve_batch)(
        lo, hi, X, X_old[lo:hi], Y, m, f, lam, lam_y, mu)
                                   for lo, hi in zip(start_idx, end_idx))
    X_new = np.vstack(res)
    return X_new