#!/usr/bin/env python
# this line is just used to define the type of document

import rospy
import mavros

from mavros_msgs.msg import OverrideRCIn
# from mavros_msgs.srv import ParamSet,ParamGet,CommandBool,SetMode

# node will publish motor speeds
from mav_msgs.msg import Actuators

#node will subscribe to odometry measurements
from nav_msgs.msg import Odometry

# controller node publishes message of this type so that GUI can plot stuff
from quad_control.msg import quad_state_and_cmd

# node will need to get state from qualysis, other sensor, or simulator
# it subscribes to a message of this tupe
from quad_control.msg import quad_state

# node will publish message quad_cmd that contains commands to simulator
from quad_control.msg import quad_cmd

# node will publish controller state if requested
from quad_control.msg import Controller_State

# import services defined in quad_control
# Four SERVICES ARE BEING USED: SaveData, ServiceTrajectoryDesired, Mocap_Id, StartSim
# SaveData is for saving data in txt file
# TrajDes is for selecting trajectory
# Mocap_Id for body detection from QUALISYS
# StartSim stop simulator
from quad_control.srv import *

# when Mocap is used this is necessary
import utilities.mocap_source

# to work with directories relative to ROS packages
from rospkg import RosPack

import numpy
from numpy import *

from utilities.utility_functions import GetRotFromEulerAnglesDeg,Velocity_Filter,Median_Filter

# import list of available trajectories
from trajectories.TrajectoryPlanner import trajectories_dictionary
from Yaw_Rate_Controller import yaw_controllers_dictionary

# from quadrotor_tracking_controllers_hierarchical import controllers_dictionary
from controllers_hierarchical.fully_actuated_controllers import controllers_dictionary

from ConverterBetweenStandards.RotorSConverter import RotorSConverter
from ConverterBetweenStandards.IrisPlusConverter import IrisPlusConverter

import math

class quad_controller():

    def __init__(self):

        # frequency of controlling action!!
        self.frequency = 35.0

        # state of quad: position, velocity and attitude
        # ROLL, PITCH, AND YAW (EULER ANGLES IN DEGREES)
        self.state_quad = numpy.zeros(3+3+3)

        # dy default, desired trajectory is staying still in origin
        TrajectoryClass    = trajectories_dictionary.trajectories_dictionary['StayAtRest']
        self.TrajGenerator = TrajectoryClass()

        # initialize counter for publishing to GUI
        # we only publish when self.PublishToGUI =1
        self.PublishToGUI = 1
        # Frequency of publishing to GUI (Hz)
        frequency_PubToGui = 10
        # we reset counter when self.PublishToGUI >= PublishToGUIBound
        self.PublishToGUIBound = int(self.frequency/frequency_PubToGui)

        # intiialization should be done in another way,
        # but median will take care of minimizing effects
        self.VelocityEstimator = Velocity_Filter(3,numpy.zeros(3),0.0)

        # controllers selected by default
        ControllerClass = controllers_dictionary.controllers_dictionary['PIDSimpleBoundedIntegralController']
        self.ControllerObject = ControllerClass()

        YawControllerClass = yaw_controllers_dictionary.yaw_controllers_dictionary['YawRateControllerTrackReferencePsi']
        self.YawControllerObject = YawControllerClass()

        # for reseting neutral value that makes iris+ stay at desired altitude
        self.DesiredZForceMedian = Median_Filter(10)

        # for saving data
        # determine ROS workspace directory
        rp = RosPack()
        # determine ROS workspace directory where data is saved
        package_path = rp.get_path('quad_control')
        self.package_save_path = package_path+'/experimental_data/data/'

    def GET_STATE_(self):

        # if simulator is on, state is updated by subscription
        # if mocap is on, state comes from mocap
        if self.flagMOCAP == True:

            self.flag_measurements = 1

            bodies = self.Qs.get_body(self.body_id)

            if bodies != 'off':

                self.flag_measurements = 1

                # position
                x=bodies["x"]
                y=bodies["y"]
                z=bodies["z"]
                p = numpy.array([x,y,z])

                # velocity
                #v = numpy.array([data.vx,data.vy,data.vz])
                v = self.VelocityEstimator.out(p,rospy.get_time())

                # attitude: euler angles THESE COME IN DEGREES
                roll = bodies["roll"]; pitch=-bodies["pitch"]; yaw  = bodies["yaw"]
                ee = numpy.array([roll,pitch,yaw])

                # collect all components of state
                self.state_quad = numpy.concatenate([p,v,ee])
            else:
                # do nothing, keep previous state
                self.flag_measurements = 2
        else:
            # simulator on
            self.flag_measurements = 0

    # callback when simulator publishes states
    def get_state_from_simulator(self, data):
        # position
        p = numpy.array([data.x,data.y,data.z])
        # velocity
        v = numpy.array([data.vx,data.vy,data.vz])
        #v = self.VelocityEstimator.out(p,rospy.get_time())
        # attitude: euler angles
        ee = numpy.array([data.roll,data.pitch,data.yaw])
        # collect all components of state
        self.state_quad = numpy.concatenate([p,v,ee])

    # callback when ROTORS simulator publishes states
    def get_state_from_rotorS_simulator(self,odometry_rotor_s):

        # # RotorS has attitude inner loop that needs to known attitude of quad
        # self.RotorSObject.rotor_s_attitude_for_control(odometry_rotor_s)
        # # get state from rotorS simulator
        # self.state_quad = self.RotorSObject.get_quad_state(odometry_rotor_s)
        return

    # for obtaining current desired state
    def traj_des(self):
        # time for trajectory generation
        time_TrajDes = rospy.get_time() - self.time_TrajDes_t0
        return self.TrajGenerator.output(time_TrajDes)

    # callback for when "saving data" is requested
    def _handle_save_data(self,req):

        if req.flag_save == True:
            # if GUI request data to be saved create file

            # namespace, e.g. /Iris1/
            namespace = rospy.get_namespace()
            if not (namespace == ""):
                # remove / symbol to namespace: e.g, we get namespace= Iris1
                namespace = namespace.replace("/", "")

            # string for time: used for generating files
            # time_stamp = str(int(rospy.get_time() - self.TimeSaveData))
            time_stamp = str(int(rospy.get_time()))

            # file name provided by user in request
            file_name    = req.file_name
            self.file_handle  = file(self.package_save_path+'_'+time_stamp+'_'+file_name+namespace+'.txt', 'w')

            # if GUI request data to be saved, set flag to true
            self.SaveDataFlag = True

        else:
            # if GUI request data NOT to be saved, set falg to False
            self.SaveDataFlag = False

        # return message to Gui, to let it know resquest has been fulfilled
        return SaveDataResponse(True)

    # callback for when "reseting iris plus neutral value" is requested
    def _handle_iris_plus_reset_neutral(self,req):
        # return message to GUI, to let it know resquest has been fulfilled
        # IrisPlusResetNeutral: service type
        median_force = self.DesiredZForceMedian.output()
        rospy.logwarn('median force = '+ str(median_force))
        self.IrisPlusConverterObject.reset_k_trottle_neutral(median_force)

        # new neutral value
        k_trottle_neutral = self.IrisPlusConverterObject.get_k_throttle_neutral()
        return IrisPlusResetNeutralResponse(received = True, k_trottle_neutral = k_trottle_neutral)

    # callback for when "seting iris plus neutral value" is requested
    def _handle_iris_plus_set_neutral(self,req):
        # return message to GUI, to let it know resquest has been fulfilled
        # IrisPlusSetNeutral: service type
        self.IrisPlusConverterObject.set_k_trottle_neutral(req.k_trottle_neutral)
        return IrisPlusSetNeutralResponse(True)

    # callback for publishing state of controller (or stop publishing)
    def handle_Controller_State_Srv(self,req):

        # if controller belongs to list of controllers that have a state
        if req.flag_controller in controllers_with_state:
            # if GUI request for state of controller, parameters = 0
            if int(req.parameters[0]) == 0:

                    # if publishging already, stop publish
                    # if not publishing, start publishing
                    self.flagPublish_ctr_st = not self.flagPublish_ctr_st

            # if GUI request for reseting state of controller, parameters = 1
            elif int(req.parameters[0]) == 1:
                self.ControllerObject.reset_estimate_xy()
            elif int(req.parameters[0]) == 2:
                self.ControllerObject.reset_estimate_z()

        # return message to Gui, to let it know resquest has been fulfilled
        return Controller_SrvResponse(True)

    # callback for when changing controller is requested
    def _handle_service_change_controller(self,req):

        # controller_class_name = req.controller_name
        # chosen class taken from dictionary
        ControllerClass = controllers_dictionary.controllers_dictionary[req.controller_name]

        self.ControllerObject = ControllerClass.from_string(req.parameters)

        #rospy.logwarn(self.ControllerObject.__class__.__name__)
        rospy.logwarn('444444444444')

        # return message to Gui, to let it know resquest has been fulfilled
        return SrvControllerChangeByStrResponse(received = True)

    # callback for when changing desired trajectory is requested
    def _handle_service_trajectory_des(self, req):

        # trajectory_class_name = req.trajectory
        # update class for TrajectoryGenerator
        TrajectoryClass = trajectories_dictionary.trajectories_dictionary[req.trajectory]

        string = req.parameters

        #offset   = numpy.array(req.offset)
        #rotation = numpy.array(req.rotation)

        self.TrajGenerator = TrajectoryClass.from_string(string)

        # we need to update initial time for trajectory generation
        self.time_TrajDes_t0 = rospy.get_time()

        # return message to Gui, to let it know resquest has been fulfilled
        return SrvTrajectoryDesiredResponse(True)

    # # function to stop simulator
    # def stop_simulator(self):
    #     try:
    #         rospy.wait_for_service('StartSimulator',1.0)

    #         try:
    #             AskForStart = rospy.ServiceProxy('StartSimulator', StartSim)

    #             reply = AskForStart(False)
    #             if reply.Started == True:
    #                 return True
    #             else:
    #                 return False
    #         except:
    #             return False

    #     except:
    #         return False

    #callback for turning ON/OFF Mocap and turning OFF/ON the subscription to the simulator
    def handle_Mocap(self,req):

        # if mocap is turned on
        if self.flagMOCAP_On == True:

            # request to turn OFF Mocap, and turn on subscription to Simulator messages
            if req.On == False:

                # in case Qs is not defined yet
                try:
                    # close mocap connection
                    # self.Qs._stop_measurement()
                    del self.Qs

                    self.flagMOCAP_On = False

                    # set flag to OFF
                    self.flagMOCAP = False

                    # subscribe again to simultor messages
                    self.SubToSim = rospy.Subscriber("quad_state", quad_state, self.get_state_from_simulator)

                    # service has been provided
                    return Mocap_IdResponse(True,True)
                except:
                    # service was NOT provided
                    return Mocap_IdResponse(True,False)

            # if we are requested to change the body Id
            if req.change_id == True:

                # see list of available bodies
                bodies = self.Qs.get_updated_bodies()

                # check if body_id available
                body_indice = -1

                # Get the corresponding id of the body
                if isinstance(bodies,list):
                    for i in range(0,len(bodies)):
                        rospy.logwarn(bodies[i]['id'])
                        if(bodies[i]['id']==req.id):
                            body_indice=i

                # save body id
                self.body_id = req.id

                if body_indice == -1:

                    # body does not exist
                    self.flagMOCAP = False

                    # body does not exist, but service was provided
                    return Mocap_IdResponse(False,True)
                else:
                    # body exists

                    # set flag to on
                    self.flagMOCAP = True

                    # body EXISTS, and service was provided
                    return Mocap_IdResponse(True,True)

                # # stop simulator
                # stop = self.stop_simulator()

                # # if simulator stopped
                # if stop == True:

                #     if body_indice == -1:

                #         # body does not exist
                #         self.flagMOCAP = False

                #         # body does not exist, but service was provided
                #         return Mocap_IdResponse(False,True)
                #     else:
                #         # body exists

                #         # set flag to on
                #         self.flagMOCAP = True

                #         # body EXISTS, and service was provided
                #         return Mocap_IdResponse(True,True)
                # else:
                #     self.flagMOCAP = False
                #     # service was NOT provided
                #     return Mocap_IdResponse(False,False)

        else:
            # if Mocap is turned off, and we are requested to turn it on
            if req.On == True:
                # establish connection to qualisys
                self.Qs = mocap_source.Mocap(info=0)

                # stop subscription to data from simulator
                # unsubscribe to topic
                self.SubToSim.unregister()

                self.flagMOCAP_On = True

                # service was provided
                return Mocap_IdResponse(False,True)

    # return list of bodies detected by mocap or numbers 1 to 99 if not available
    def handle_available_bodies(self, dummy):
        if self.flagMOCAP_On:
            try: # sometimes mocap causes unpredictable errors
                bodies = self.Qs.find_available_bodies(False)
                if len(bodies) > 0:
                    return {"bodies": bodies[0]}
            except:
                pass

        return {"bodies": range(0,100)}

    def PublishToGui(self,states_d,Input_to_Quad):

        # WE ONLY PUBLIS TO TO GUI AT A CERTAIN FREQEUNCY
        # WHICH IS NOT NECESSARILY THE FREQUENCY OF THE NODE
        if self.PublishToGUI <= self.PublishToGUIBound:
            # if we dont publish, we increase counter
            self.PublishToGUI = self.PublishToGUI + 1
        else:
            # if we publish, we increase counter
            self.PublishToGUI = 1

            # create a message of type quad_state_and_cmd
            st_cmd = quad_state_and_cmd()

            # get current time
            st_cmd.time  = rospy.get_time()

            # state of quad comes from QUALISYS, or other sensor
            st_cmd.x     = self.state_quad[0]; st_cmd.y     = self.state_quad[1]; st_cmd.z     = self.state_quad[2]
            st_cmd.vx    = self.state_quad[3]; st_cmd.vy    = self.state_quad[4]; st_cmd.vz    = self.state_quad[5]
            st_cmd.roll  = self.state_quad[6]; st_cmd.pitch = self.state_quad[7]; st_cmd.yaw   = self.state_quad[8]

            st_cmd.xd    = states_d[0]; st_cmd.yd    = states_d[1]; st_cmd.zd    = states_d[2]
            st_cmd.vxd   = states_d[3]; st_cmd.vyd   = states_d[4]; st_cmd.vzd   = states_d[5]

            st_cmd.cmd_1 = Input_to_Quad[0]; st_cmd.cmd_2 = Input_to_Quad[1]; st_cmd.cmd_3 = Input_to_Quad[2]; st_cmd.cmd_4 = Input_to_Quad[3]

            st_cmd.cmd_5 = 1500.0; st_cmd.cmd_6 = 1500.0; st_cmd.cmd_7 = 1500.0; st_cmd.cmd_8 = 1500.0

            self.pub.publish(st_cmd)

            # controller state is supposed to be published
            if self.flagPublish_ctr_st:
                # publish controller state
                msg       = Controller_State()
                msg.time  = rospy.get_time()
                msg.d_est = self.ControllerObject.d_est
                self.pub_ctr_st.publish(msg)

    def PublishToQuad(self,Input_to_Quad):

        # create a message of the type quad_cmd, that the simulator subscribes to
        cmd  = quad_cmd()

        cmd.cmd_1 = Input_to_Quad[0]; cmd.cmd_2 = Input_to_Quad[1]; cmd.cmd_3 = Input_to_Quad[2]; cmd.cmd_4 = Input_to_Quad[3];

        cmd.cmd_5 = 1500.0; cmd.cmd_6 = 1500.0; cmd.cmd_7 = 1500.0; cmd.cmd_8 = 1500.0

        self.pub_cmd.publish(cmd)

    def control_compute(self):

        # node will be named quad_control (see rqt_graph)
        rospy.init_node('quad_control', anonymous=True)

        # message published by quad_control to GUI
        self.pub = rospy.Publisher('quad_state_and_cmd', quad_state_and_cmd, queue_size=10)

        # message published by quad_control that simulator will subscribe to
        self.pub_cmd = rospy.Publisher('quad_cmd', quad_cmd, queue_size=10)

        # for publishing state of the controller
        self.pub_ctr_st = rospy.Publisher('ctr_state', Controller_State, queue_size=10)
        # initialize flag for publishing controller state at false
        self.flagPublish_ctr_st = False

        # controller needs to have access to STATE of the system
        # this can come from QUALISYS, a sensor, or the simulator
        self.SubToSim = rospy.Subscriber("quad_state", quad_state, self.get_state_from_simulator)

        #-----------------------------------------------------------------------#
        # TO SAVE DATA FLAG
        # by default, NO data is saved
        self.SaveDataFlag = False
        # we will use this just for convenience, when generating the names of the files
        # where the data will be saved
        self.TimeSaveData = rospy.get_time()
        # Service is created, so that data is saved when GUI requests
        Save_data_service = rospy.Service('SaveDataFromGui', SaveData, self._handle_save_data)

        #-----------------------------------------------------------------------#
        # SERVICE FOR SELECTING DESIRED TRAJECTORY
        # by default, STAYING STILL IN ORIGIN IS DESIRED TRAJECTORY
        # self.flagTrajDes = 0
        # Service is created, so that data is saved when GUI requests
        TrajDes_service = rospy.Service('ServiceTrajectoryDesired', SrvTrajectoryDesired, self._handle_service_trajectory_des)

        # initialize initial time for trajectory generation
        self.time_TrajDes_t0 = rospy.get_time()

        #-----------------------------------------------------------------------#
        # flag for MOCAP is initialized as FALSE
        # flag for wheter Mocap is being used or not
        self.flagMOCAP    = False
        # Flag for whether Mocap is ON or OFF
        self.flagMOCAP_On = False
        # Service is created, so that Mocap is turned ON or OFF whenever we want
        Save_MOCAP_service = rospy.Service('Mocap_Set_Id', Mocap_Id, self.handle_Mocap)

        # Service for providing list of available mocap bodies to GUI
        mocap_available_bodies = rospy.Service('MocapBodies', MocapBodies, self.handle_available_bodies)

        #-----------------------------------------------------------------------#
        # Service is created, so that user can change controller on GUI
        rospy.Service('ServiceChangeController', SrvControllerChangeByStr, self._handle_service_change_controller)
        # rospy.Service('ServiceChangeController', SrvControllerChange, self._handle_service_change_controller)

        #-----------------------------------------------------------------------#
        # Service: change neutral value that guarantees that a quad remains at a desired altitude
        rospy.Service('IrisPlusResetNeutral', IrisPlusResetNeutral, self._handle_iris_plus_reset_neutral)
        rospy.Service('IrisPlusSetNeutral', IrisPlusSetNeutral, self._handle_iris_plus_set_neutral)

        #-----------------------------------------------------------------------#
        # Service for publishing state of controller
        # we use same type of service as above
        #Contller_St = rospy.Service('Controller_State_GUI', Controller_Srv, self.handle_Controller_State_Srv)

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

        self.RotorSObject = RotorSConverter()
        self.RotorSFlag   = True

        # subscriber: to odometry
        # self.sub_odometry = rospy.Subscriber("/firefly/odometry_sensor1/odometry", Odometry, self.compute_cmd)
        self.sub_odometry = rospy.Subscriber("/firefly/ground_truth/odometry", Odometry, self.get_state_from_rotorS_simulator)

        # # subscriber: to odometry
        # # self.sub_odometry = rospy.Subscriber("/firefly/odometry_sensor1/odometry", Odometry, self.compute_cmd)
        # self.sub_odometry_load = rospy.Subscriber("/firefly/ground_truth/odometry_load", Odometry, self.update_load_odometry)

        # publisher: command firefly motor speeds
        self.pub_motor_speeds = rospy.Publisher('/firefly/command/motor_speed', Actuators, queue_size=10)
        #-----------------------------------------------------------------------#

        # rc_override = rospy.Publisher('mavros/rc/override',OverrideRCIn,queue_size=10)
        rc_override = rospy.Publisher('mavros/rc/override', OverrideRCIn, queue_size=100)

        pub = rospy.Publisher('mavros/rc/override', OverrideRCIn, queue_size=100)

        rate = rospy.Rate(self.frequency)

        t0 = rospy.get_time()

        self.IrisPlusConverterObject = IrisPlusConverter()

        while not rospy.is_shutdown():

            time = rospy.get_time()

            # get state: if MOCAP is ON we ask MOCAP; if MOCAP in OFF, we are subscribed to Simulator topic
            self.GET_STATE_()

            #print(self.state_quad[0:3])
            # print self.state_quad[6:9]

            # states for desired trajectory
            states_d = self.traj_des()

            # compute input to send to QUAD
            # rospy.logwarn(self.ControllerObject.__class__.__name__)
            # rospy.logwarn(self.state_quad[0:3])
            # rospy.logwarn('bbbbbbbbbbbbbbbbbbbbbb')
            # rospy.logwarn(states_d[0:3])
            # rospy.logwarn('ccccccccccccccccccccccccccc')
            desired_3d_force_quad = self.ControllerObject.output(time,self.state_quad,states_d)
            # rospy.logwarn('dddddddddddddd')
            # rospy.logwarn(desired_3d_force_quad)
            self.DesiredZForceMedian.update_data(desired_3d_force_quad[2])

            euler_rad     = self.state_quad[6:9]*math.pi/180
            euler_rad_dot = numpy.zeros(3)

            # convert to iris + standard
            state_for_yaw_controller = numpy.concatenate([euler_rad,euler_rad_dot])
            input_for_yaw_controller = numpy.array([0.0,0.0])
            yaw_rate = self.YawControllerObject.output(state_for_yaw_controller,input_for_yaw_controller)

            self.IrisPlusConverterObject.set_rotation_matrix(euler_rad)
            iris_plus_rc_input = self.IrisPlusConverterObject.input_conveter(desired_3d_force_quad,yaw_rate)

            # Publish commands to Quad
            self.PublishToQuad(iris_plus_rc_input)

            # publish to GUI (it also contains publish state of Control to GUI)
            self.PublishToGui(states_d,iris_plus_rc_input)

            # ORDER OF INPUTS IS VERY IMPORTANT: roll, pitch, throttle,yaw_rate
            # create message of type OverrideRCIn
            rc_cmd          = OverrideRCIn()
            other_channels  = numpy.array([1500,1500,1500,1500])
            channels        = numpy.concatenate([iris_plus_rc_input,other_channels])
            channels        = numpy.array(channels,dtype=numpy.uint16)
            rc_cmd.channels = channels
            rc_override.publish(rc_cmd)

            # publish message
            # TOTAL_MASS = 1.66779; Force3D = numpy.array([0.0,0.0,TOTAL_MASS*9.85]); PsiStar = 0.0; self.pub_motor_speeds.publish(self.RotorSObject.rotor_s_message(Force3D,PsiStar))
            self.pub_motor_speeds.publish(self.RotorSObject.rotor_s_message(desired_3d_force_quad,yaw_rate))

            if self.SaveDataFlag == True:
                # if we want to save data
                # numpy.savetxt(self.file_handle, [concatenate([[rospy.get_time()],[self.flag_measurements], self.state_quad, states_d[0:9], Input_to_Quad,self.ControllerObject.d_est])],delimiter=' ')
                numpy.savetxt(self.file_handle, [concatenate([[rospy.get_time()],[self.flag_measurements], self.state_quad, states_d[0:9], desired_3d_force_quad])],delimiter=' ')
            # go to sleep
            rate.sleep()

if __name__ == '__main__':
    A_Controller = quad_controller()
    try:
        A_Controller.control_compute()
    except rospy.ROSInterruptException:
        pass


import asyncio
import time

async def fun1():
    print("func1")
    await asyncio.sleep(3)
    print("fun1")
    return 'fun1--result'

async def fun2():
    print("func2")
    await asyncio.sleep(2)
    print("fun2")
    return 'fun2--result'

async def fun3():
    print("func3")
    await asyncio.sleep(1)
    print("fun3")
    return 'fun3--result'

async def main():
    task = [
        asyncio.create_task(fun1()),
        asyncio.create_task(fun2()),
        asyncio.create_task(fun3())
    ]
    await asyncio.wait(task)

if __name__ == '__main__':
    start = time.time()
    loop = asyncio.get_event_loop()
    # task1 = loop.create_task(fun1())
    # task2 = loop.create_task(fun2())
    # task3 = loop.create_task(fun3())
    loop.run_until_complete(main())
    # asyncio.run(main())
    # asyncio.run(asyncio.wait([fun1(), fun2(), fun3()])) # 同时执行多个任务
    end = time.time()
    print(end - start)


import argparse

parser = argparse.ArgumentParser()

parser.add_argument("values", nargs='*')
parser.add_argument("--sort", action="store_true")

args = parser.parse_args()

lst = args.values

d = dict()

for string in lst:
    key, val = string.split('=')
    d[key] = val

keys = list(d.keys())
if args.sort:
    keys.sort()

for key in keys:
    print(f"Key: {key} Value: {d[key]}")


from arcade.gui.widgets import UIDummy, UIBorder

def test_wrap_with_default_border():
    # GIVEN
    widget = UIDummy()

    # WHEN
    with_border = widget.with_border()

    # THEN
    assert isinstance(with_border, UIBorder)
    assert with_border.rect == (-2, -2, 104, 104)

def test_wrap_with_extra_wide_border():
    # GIVEN
    border_width = 10
    widget = UIDummy()

    # WHEN
    with_border = widget.with_border(width=border_width)

    # THEN
    assert isinstance(with_border, UIBorder)
    assert with_border.rect == (-border_width, -border_width, 100 + 2 * border_width, 100 + 2 * border_width)


import skimage
from skimage import io, measure
from sklearn import cluster
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math

parrots = io.imread('parrots.jpg')
parrots_float = skimage.img_as_float(parrots)
x, y, z = parrots_float.shape
obj_feat_matrix = parrots_float.reshape(x * y, z)

for j in range(5, 20):

    km = cluster.KMeans(random_state=241, n_clusters=j)
    km.fit(obj_feat_matrix)

    df = pd.DataFrame(obj_feat_matrix, columns=list('RGB'))
    df['Cluster'] = km.predict(df)

    df_mean_groups = df.groupby('Cluster').mean()
    df_median_groups = df.groupby('Cluster').median()
    df_mean = df.copy()
    for i in range(j):
        df_mean.ix[df_mean['Cluster'] == i, [0, 1, 2]] = df_mean_groups.ix[i].tolist()

    df_median = df.copy()
    for i in range(j):
        df_median.ix[df_median['Cluster'] == i, [0, 1, 2]] = df_median_groups.ix[i].tolist()

    MSE = 0

    MSE_mean = ((df.ix[:, 0:3] - df_mean.ix[:, 0:3]) ** 2).mean()
    MSE_mean = sum(MSE_mean) / 3
    PSNR_mean = 10 * math.log(pow(1, 2) / MSE_mean, 10)
    print("PSNR mean:", PSNR_mean, "j:", j)

    MSE_median = ((df.ix[:, 0:3] - df_median.ix[:, 0:3]) ** 2).median()
    MSE_median = sum(MSE_median) / 3
    PSNR_median = 10 * math.log(pow(1, 2) / MSE_median, 10)
    print("PSNR median:", PSNR_median, "j:", j)

    print()


import logging

from core.utils import twitter
from core.models import Profile
from core.models import Settings

from django.conf import settings

logger = logging.getLogger('core.cron')

def sync():
    profile = Profile.objects.filter(user__username='5833511420').first()

    if profile is None:
        return

    last_tweet_id = Settings.last_tweet_id()

    timeline = twitter.get_timeline(since_id=last_tweet_id.value)
    timeline.reverse()

    oauth_weibo = settings.WEIBO

    count = 0
    for tweet in timeline:
        weibo = tweet.to_weibo()
        if weibo is not None:
            result = oauth_weibo.post(profile, weibo)

            if result:
                last_tweet_id.value = weibo.tweet_id
                last_tweet_id.save()
                count += 1

    logger.info(f'Cron job finished. Length: {len(timeline)}, sent: {count}')


import pygame, sys
from pygame.locals import *

def key_down_events(event, player, sprite_animation, setting, camera, enemy, menu):
    if event.key == K_RIGHT:
        player.x_change = setting.player_speed
        player.right = True
        sprite_animation.stop_left = False
    elif event.key == K_LEFT:
        player.x_change = -setting.player_speed
        player.left = True
        sprite_animation.stop_left = True
    elif event.key == K_UP:
        menu.up_arrow = True
        player.y_change = setting.jump_speed
        player.jump = True
    elif event.key == K_DOWN:
        menu.down_arrow = True
    elif event.key == K_RETURN:
        if menu.menu_select_rect.y != menu.credits_rect.centery - 2:
            setting.start_game = True

    elif event.key == K_q:
        pygame.quit()
        sys.exit()

def key_up_events(event, player, sprite_animation, menu):
     if event.key == K_RIGHT:
        player.right = False
        if player.left == False:
            player.x_change = 0
     elif event.key == K_LEFT:
        player.left = False
        if player.right == False:
            player.x_change = 0
     elif event.key == K_UP:
        player.jump = False
        menu.up_arrow = False
     elif event.key == K_DOWN:
        menu.down_arrow = False

def game_events(player, sprite_animation, setting, camera, enemy, menu):
    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()
        elif event.type == KEYDOWN:
            key_down_events(event, player, sprite_animation, setting, camera, enemy, menu)
        elif event.type == KEYUP:
            key_up_events(event, player, sprite_animation, menu)

def update_screen(player, sprite_animation, setting, camera, enemy, menu, tile_map, death, moving_objects, coin):
    if setting.start_game:
        if enemy.collide_enemy == False: # make player visible, not dead
                player.update()
                sprite_animation.update()
        moving_objects.update()
        camera.update_camera(sprite_animation, enemy)
        tile_map.blit_map(camera, player)
        enemy.update()
        death.update()
        coin.update(camera)
    if setting.start_game == False:
        menu.update()
        reset_game(setting, camera, player, sprite_animation, enemy, moving_objects, tile_map, coin)

"""reset back to initial values"""
def reset_game(setting, camera, player, sprite_animation, enemy, moving_objects, tile_map, coin):
    setting.start_game = False
    sprite_animation.stop_left = False
    enemy.collide_enemy = False
    player.collide_box_item = False
    camera.camera_position = 0
    player.rect.x = 400
    player.rect.bottom = 750
    moving_objects.rect.x = 2000 # Asteroid 1
    moving_objects.rect_2.x = 4000 # Asteroid
    tile_map.tile_map_level[2][30] = setting.box_item
    tile_map.tile_map_level[4][17] = setting.box_item
    tile_map.tile_map_level[10][27] = setting.box_item
    tile_map.tile_map_level[10][43] = setting.box_item
    coin.coin1 = True
    coin.coin2 = True
    coin.coin3 = True
    coin.coin4 = True
    coin.coins_collected = 0
    tile_map.show_coin1 = False
    tile_map.show_coin2 = False
    tile_map.show_coin3 = False
    tile_map.show_coin4 = False


"""
Experiment: dataset cleaning from noisy labels,
            specifically a percentage of flipped labels
            in the train data. Only for binary classification datasets.
"""
import time
import uuid
import argparse
from datetime import datetime
from copy import deepcopy
import os
import sys
import json
import warnings
warnings.simplefilter(action='ignore', category=UserWarning)  # lgb compiler warning
here = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, here + '/../../')  # for influence_boosting
sys.path.insert(0, here + '/../')  # for utility

import tqdm
import numpy as np
from sklearn.base import clone
from sklearn.metrics import accuracy_score
from sklearn.metrics.pairwise import rbf_kernel
from sklearn.neighbors import KNeighborsClassifier

import trex
from utility import model_util
from utility import data_util
from utility import exp_util
from utility import print_util
from baselines.influence_boosting.influence.leaf_influence import CBLeafInfluenceEnsemble
from baselines.maple.MAPLE import MAPLE
from baselines.mmd_critic import mmd

TREE_PROTO_K = 10
N_PROTOTYPES = 10

def _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy):
    """
    Retrains the tree ensemble for each ckeckpoint, where a checkpoint represents
    which flipped labels have been fixed.
    """

    X_train, y_train, X_test, y_test = data

    accs = [acc_test_noisy]
    checked_pct = [0]
    fix_pct = [0]

    for n_checked, n_fix in tqdm.tqdm(ckpt_ndx):
        fix_indices = fix_ndx[:n_fix]

        semi_noisy_ndx = np.setdiff1d(noisy_ndx, fix_indices)
        y_train_semi_noisy = data_util.flip_labels_with_indices(y_train, semi_noisy_ndx)

        model_semi_noisy = clone(clf).fit(X_train, y_train_semi_noisy)
        acc_test = accuracy_score(y_test, model_semi_noisy.predict(X_test))

        accs.append(acc_test)
        checked_pct.append(float(n_checked / len(y_train)))
        fix_pct.append(float(n_fix / len(noisy_ndx)))

    return checked_pct, accs

def _record_fixes(train_order, noisy_ndx, train_len, interval):
    """
    Returns the number of train instances checked and which train instances were
    fixed for each checkpoint.
    """

    fix_ndx = []
    ckpt_ndx = []
    checked = 0
    snapshot = 1

    for train_ndx in train_order:
        if train_ndx in noisy_ndx:
            fix_ndx.append(train_ndx)
        checked += 1

        if float(checked / train_len) >= (snapshot * interval):
            ckpt_ndx.append((checked, len(fix_ndx)))
            snapshot += 1
    fix_ndx = np.array(fix_ndx)

    return ckpt_ndx, fix_ndx

def _our_method(explainer, noisy_ndx, y_train, n_check, interval):
    """
    Order instnces by largest absolute values of TREX weights.
    """
    train_weight = explainer.get_weight()[0]
    train_order = np.argsort(np.abs(train_weight))[::-1][:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, len(y_train), interval)

    return ckpt_ndx, fix_ndx, train_weight

def _random_method(noisy_ndx, y_train, interval, to_check=1, random_state=1):
    """
    Randomly picks train instances from the train data.
    """

    n_train = len(y_train)

    assert to_check <= n_train, 'to_check > n_train!'
    if isinstance(to_check, int):
        n_check = to_check
    else:
        exit('to_check not int')

    np.random.seed(random_state + 1)  # +1 to avoid choosing the same indices as the noisy labels
    train_order = np.random.choice(n_train, size=n_check, replace=False)
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, n_train, interval)
    return ckpt_ndx, fix_ndx

def _loss_method(noisy_ndx, y_train_proba, y_train, interval, to_check=1, logloss=False):
    """
    Sorts train instances by largest train loss.
    """

    n_train = len(y_train)

    assert to_check <= n_train, 'to_check > n_train!'
    if isinstance(to_check, int):
        n_check = to_check
    else:
        exit('to_check not int')

    # extract 1d array of probabilities representing the probability of the target label
    if y_train_proba.ndim > 1:
        y_proba = model_util.positive_class_proba(y_train, y_train_proba)
    else:
        y_proba = y_train_proba

    # compute the loss for each instance
    y_loss = exp_util.instance_loss(y_proba, y_train, logloss=logloss)

    # put train instances in order based on decreasing absolute loss
    if logloss:
        train_order = np.argsort(y_loss)[:n_check]  # ascending order, most negative first
    else:
        train_order = np.argsort(y_loss)[::-1][:n_check]  # descending order, most positive first

    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, len(y_train), interval)
    return ckpt_ndx, fix_ndx, y_loss, train_order

def _influence_method(explainer, noisy_ndx, X_train, y_train, y_train_noisy, interval, to_check=1):
    """
    Computes the influence on train instance i if train instance i were upweighted/removed.
    Reference: https://github.com/kohpangwei/influence-release/blob/master/influence/experiments.py
    This uses the fastleafinfluence method by Sharchilev et al.
    """

    n_train = len(y_train)

    assert to_check <= n_train, 'to_check > n_train!'
    if isinstance(to_check, int):
        n_check = to_check
    else:
        exit('to_check not int')

    influence_scores = []
    buf = deepcopy(explainer)
    for i in tqdm.tqdm(range(len(X_train))):
        explainer.fit(removed_point_idx=i, destination_model=buf)
        influence_scores.append(buf.loss_derivative(X_train[[i]], y_train_noisy[[i]])[0])
    influence_scores = np.array(influence_scores)

    # sort by ascending order; the most negative train instances
    # are the ones that increase the log loss the most, and are the most harmful
    train_order = np.argsort(influence_scores)[:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, n_train, interval)
    return ckpt_ndx, fix_ndx, influence_scores, train_order

def _maple_method(explainer, X_train, noisy_ndx, interval, to_check=1):
    """
    Orders instances by tree kernel similarity density.
    """

    n_train = X_train.shape[0]

    assert to_check <= n_train, 'to_check > n_train!'
    if isinstance(to_check, int):
        n_check = to_check
    else:
        exit('to_check not int')

    train_weight = np.zeros(X_train.shape[0])
    for i in tqdm.tqdm(range(X_train.shape[0])):
        weights = explainer.get_weights(X_train[i])
        train_weight[i] += np.sum(weights)

    train_order = np.argsort(train_weight)[::-1][:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, n_train, interval)
    return ckpt_ndx, fix_ndx, train_weight, train_order

def _knn_method(knn_clf, X_train, noisy_ndx, interval, to_check=1):
    """
    Count impact by the number of times a training sample shows up in
    one another's neighborhoods, this can be weighted by 1 / distance.
    """

    n_train = X_train.shape[0]

    assert to_check <= n_train, 'to_check > n_train!'
    if isinstance(to_check, int):
        n_check = to_check
    else:
        exit('to_check not int')

    train_impact = np.zeros(X_train.shape[0])
    for i in tqdm.tqdm(range(X_train.shape[0])):
        distances, neighbor_ids = knn_clf.kneighbors([X_train[i]])
        train_impact[neighbor_ids[0]] += 1

    train_order = np.argsort(train_impact)[::-1][:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, n_train, interval)
    return ckpt_ndx, fix_ndx, train_order

def _proto_method(model, X_train, y_train, noisy_ndx, interval, n_check):
    """
    Orders instances by using the GBT distance similarity formula in
    https://arxiv.org/pdf/1611.07115.pdf, then ranks training samples
    based on the proportion of labels from the k = 10, nearest neighbors.
    """
    extractor = trex.TreeExtractor(model, tree_kernel='leaf_path')
    X_train_alt = extractor.fit_transform(X_train)

    # obtain weight of each tree: note, this code is specific to CatBoost
    temp_fp = '.{}_cb.json'.format(str(uuid.uuid4()))
    model.save_model(temp_fp, format='json')
    cb_dump = json.load(open(temp_fp, 'r'))

    # obtain weight of each tree: learning_rate^2 * var(predictions)
    tree_weights = []
    for tree in cb_dump['oblivious_trees']:
        predictions = []

        for val, weight in zip(tree['leaf_values'], tree['leaf_weights']):
            predictions += [val] * weight

        tree_weights.append(np.var(predictions) * (model.learning_rate_ ** 2))

    # weight leaf path feature representation by the tree weights
    for i in range(X_train_alt.shape[0]):

        weight_cnt = 0
        for j in range(X_train_alt.shape[1]):

            if X_train_alt[i][j] == 1:
                X_train_alt[i][j] *= tree_weights[weight_cnt]
                weight_cnt += 1

        assert weight_cnt == len(tree_weights)

    # build a KNN using this proximity measure using k = 10
    knn = KNeighborsClassifier(n_neighbors=TREE_PROTO_K)
    knn = knn.fit(X_train_alt, y_train)

    # compute proportion of neighbors that share the same label
    train_impact = np.zeros(X_train_alt.shape[0])
    for i in tqdm.tqdm(range(X_train_alt.shape[0])):
        _, neighbor_ids = knn.kneighbors([X_train_alt[i]])
        train_impact[i] = len(np.where(y_train[i] == y_train[neighbor_ids[0]])[0]) / len(neighbor_ids[0])

    # rank training instances by low label agreement with its neighbors
    train_order = np.argsort(train_impact)[:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, len(y_train), interval)

    os.system('rm {}'.format(temp_fp))

    return ckpt_ndx, fix_ndx

def _mmd_method(model, X_train, y_train, noisy_ndx, interval, n_check):
    """
    Orders instances by prototypes and/or criticisms.
    """
    n_prototypes = N_PROTOTYPES
    n_criticisms = n_check

    X = np.hstack([X_train, y_train.reshape(-1, 1)])
    K = rbf_kernel(X)
    candidates = np.arange(len(K))

    prototypes = mmd.greedy_select_protos(K, candidates, n_prototypes)
    criticisms = mmd.select_criticism_regularized(K, prototypes, n_criticisms, is_K_sparse=False)

    train_order = criticisms[:n_check]
    ckpt_ndx, fix_ndx = _record_fixes(train_order, noisy_ndx, len(y_train), interval)

    return ckpt_ndx, fix_ndx

def experiment(args, logger, out_dir, seed):
    """
    Main method that trains a tree ensemble, flips a percentage of train labels, prioritizes train
    instances using various methods, and computes how effective each method is at cleaning the data.
    """

    # get model and data
    clf = model_util.get_classifier(args.tree_type,
                                    n_estimators=args.n_estimators,
                                    max_depth=args.max_depth,
                                    random_state=seed)

    X_train, X_test, y_train, y_test, label = data_util.get_data(args.dataset,
                                                                 random_state=seed,
                                                                 data_dir=args.data_dir)

    # reduce train size
    if args.train_frac < 1.0 and args.train_frac > 0.0:
        n_train = int(X_train.shape[0] * args.train_frac)
        X_train, y_train = X_train[:n_train], y_train[:n_train]
    data = X_train, y_train, X_test, y_test

    logger.info('no. train instances: {:,}'.format(len(X_train)))
    logger.info('no. test instances: {:,}'.format(len(X_test)))
    logger.info('no. features: {:,}'.format(X_train.shape[1]))

    # add noise
    y_train_noisy, noisy_ndx = data_util.flip_labels(y_train, k=args.flip_frac, random_state=seed)
    noisy_ndx = np.array(sorted(noisy_ndx))
    logger.info('no. noisy labels: {:,}'.format(len(noisy_ndx)))

    # train a tree ensemble on the clean and noisy labels
    model = clone(clf).fit(X_train, y_train)
    model_noisy = clone(clf).fit(X_train, y_train_noisy)

    # show model performance before and after noise
    logger.info('\nBefore noise:')
    model_util.performance(model, X_train, y_train, X_test=X_test, y_test=y_test, logger=logger)
    logger.info('\nAfter noise:')
    model_util.performance(model_noisy, X_train, y_train_noisy, X_test=X_test, y_test=y_test, logger=logger)

    # check accuracy before and after noise
    acc_test_clean = accuracy_score(y_test, model.predict(X_test))
    acc_test_noisy = accuracy_score(y_test, model_noisy.predict(X_test))

    # find how many corrupted/non-corrupted labels were incorrectly predicted
    if not args.true_label:
        logger.info('\nUsing predicted labels:')
        predicted_labels = model_noisy.predict(X_train).flatten()
        incorrect_ndx = np.where(y_train_noisy != predicted_labels)[0]
        incorrect_corrupted_ndx = np.intersect1d(noisy_ndx, incorrect_ndx)
        logger.info('incorrectly predicted corrupted labels: {:,}'.format(incorrect_corrupted_ndx.shape[0]))
        logger.info('total number of incorrectly predicted labels: {:,}'.format(incorrect_ndx.shape[0]))

    # number of checkpoints to record
    n_check = int(len(y_train) * args.check_pct)
    interval = (n_check / len(y_train)) / args.n_plot_points

    # random method
    logger.info('\nordering by random...')
    start = time.time()
    ckpt_ndx, fix_ndx = _random_method(noisy_ndx, y_train, interval,
                                       to_check=n_check,
                                       random_state=seed)
    check_pct, random_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)
    logger.info('time: {:3f}s'.format(time.time() - start))
    np.save(os.path.join(out_dir, 'random.npy'), random_res)

    # save global lines
    np.save(os.path.join(out_dir, 'test_clean.npy'), acc_test_clean)
    np.save(os.path.join(out_dir, 'check_pct.npy'), check_pct)

    # tree loss method
    logger.info('\nordering by tree loss...')
    start = time.time()

    y_train_proba = model_noisy.predict_proba(X_train)
    ckpt_ndx, fix_ndx, _, _ = _loss_method(noisy_ndx, y_train_proba, y_train_noisy, interval, to_check=n_check)
    _, tree_loss_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

    logger.info('time: {:3f}s'.format(time.time() - start))
    np.save(os.path.join(out_dir, 'tree.npy'), tree_loss_res)

    # trex method
    if args.trex:
        logger.info('\nordering by TREX...')
        start = time.time()
        explainer = trex.TreeExplainer(model_noisy, X_train, y_train_noisy,
                                       tree_kernel=args.tree_kernel,
                                       random_state=seed,
                                       true_label=args.true_label,
                                       kernel_model=args.kernel_model,
                                       verbose=args.verbose,
                                       val_frac=args.val_frac,
                                       logger=logger)

        ckpt_ndx, fix_ndx, _ = _our_method(explainer, noisy_ndx, y_train, n_check, interval)
        check_pct, trex_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), trex_res)

        # trex loss method
        logger.info('\nordering by TREX loss...')
        start = time.time()

        y_train_proba = explainer.predict_proba(X_train)
        ckpt_ndx, fix_ndx, _, _ = _loss_method(noisy_ndx, y_train_proba, y_train_noisy, interval, to_check=n_check)
        _, trex_loss_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method_loss.npy'), trex_loss_res)

    # influence method
    if args.tree_type == 'cb' and args.inf_k is not None:
        logger.info('\nordering by leafinfluence...')
        start = time.time()

        model_path = '.model.json'
        model_noisy.save_model(model_path, format='json')

        if args.inf_k == -1:
            update_set = 'AllPoints'
        elif args.inf_k == 0:
            update_set = 'SinglePoint'
        else:
            update_set = 'TopKLeaves'

        leaf_influence = CBLeafInfluenceEnsemble(model_path, X_train, y_train_noisy, k=args.inf_k,
                                                 learning_rate=model.learning_rate_, update_set=update_set)
        ckpt_ndx, fix_ndx, _, _ = _influence_method(leaf_influence, noisy_ndx, X_train, y_train, y_train_noisy,
                                                    interval, to_check=n_check)
        _, leafinfluence_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), leafinfluence_res)

    # MAPLE method
    if args.maple:
        logger.info('\nordering by MAPLE...')
        start = time.time()

        train_label = y_train_noisy if args.true_label else model_noisy.predict(X_train)
        maple_exp = MAPLE(X_train, train_label, X_train, train_label, verbose=args.verbose, dstump=False)
        ckpt_ndx, fix_ndx, map_scores, map_order = _maple_method(maple_exp, X_train, noisy_ndx, interval,
                                                                 to_check=n_check)
        _, maple_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), maple_res)

    # TEKNN method
    if args.teknn:
        logger.info('\nordering by teknn...')
        start = time.time()

        # transform the data
        extractor = trex.TreeExtractor(model_noisy, tree_kernel=args.tree_kernel)
        X_train_alt = extractor.fit_transform(X_train)
        train_label = y_train if args.true_label else model_noisy.predict(X_train)

        # tune and train teknn
        knn_clf = exp_util.tune_knn(model_noisy, X_train, X_train_alt, train_label, args.val_frac,
                                    seed=seed, logger=logger)

        ckpt_ndx, fix_ndx, _ = _knn_method(knn_clf, X_train_alt, noisy_ndx, interval, to_check=n_check)
        _, teknn_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), teknn_res)

        # TEKNN loss method
        logger.info('\nordering by teknn loss...')
        start = time.time()
        y_train_proba = knn_clf.predict_proba(X_train_alt)

        ckpt_ndx, fix_ndx, _, _ = _loss_method(noisy_ndx, y_train_proba, y_train_noisy, interval, to_check=n_check)
        _, teknn_loss_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method_loss.npy'), teknn_loss_res)

    # MMD-Critic method
    if args.mmd:
        logger.info('\nordering by mmd-critic...')
        start = time.time()
        ckpt_ndx, fix_ndx = _mmd_method(model_noisy, X_train, y_train_noisy, noisy_ndx, interval, n_check)
        _, mmd_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), mmd_res)

    # Prototype method
    if args.proto:
        logger.info('\nordering by proto...')
        start = time.time()
        ckpt_ndx, fix_ndx = _proto_method(model_noisy, X_train, y_train_noisy, noisy_ndx, interval, n_check)
        _, proto_res = _interval_performance(ckpt_ndx, fix_ndx, noisy_ndx, clf, data, acc_test_noisy)

        logger.info('time: {:3f}s'.format(time.time() - start))
        np.save(os.path.join(out_dir, 'method.npy'), proto_res)

def main(args):

    # make logger
    dataset = args.dataset

    if args.train_frac < 1.0 and args.train_frac > 0.0:
        dataset += '_{}'.format(str(args.train_frac).replace('.', 'p'))

    out_dir = os.path.join(args.out_dir, dataset, args.tree_type,
                           'rs{}'.format(args.rs))

    if args.trex:
        out_dir = os.path.join(out_dir, args.kernel_model, args.tree_kernel)
    elif args.teknn:
        out_dir = os.path.join(out_dir, 'teknn', args.tree_kernel)
    elif args.maple:
        out_dir = os.path.join(out_dir, 'maple')
    elif args.inf_k is not None:
        out_dir = os.path.join(out_dir, 'leaf_influence')
    elif args.mmd:
        out_dir = os.path.join(out_dir, 'mmd')
    elif args.proto:
        out_dir = os.path.join(out_dir, 'proto')

    os.makedirs(out_dir, exist_ok=True)
    logger = print_util.get_logger(os.path.join(out_dir, 'log.txt'))
    logger.info(args)
    logger.info(datetime.now())

    seed = args.rs
    logger.info('\nSeed: {}'.format(seed))
    experiment(args, logger, out_dir, seed=seed)
    print_util.remove_logger(logger)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Feature representation extractions for tree ensembles',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # I/O settings
    parser.add_argument('--dataset', type=str, default='adult', help='dataset to explain.')
    parser.add_argument('--data_dir', type=str, default='data', help='data directory.')
    parser.add_argument('--out_dir', type=str, default='output/cleaning/', help='output directory.')

    # data settings
    parser.add_argument('--train_frac', type=float, default=1.0, help='amount of training data to use.')
    parser.add_argument('--val_frac', type=float, default=0.1, help='amount of training data to use for validation.')
    parser.add_argument('--flip_frac', type=float, default=0.4, help='fraction of train labels to flip.')

    # tree settings
    parser.add_argument('--tree_type', type=str, default='cb', help='tree model to use.')
    parser.add_argument('--n_estimators', type=int, default=100, help='number of trees in random forest.')
    parser.add_argument('--max_depth', type=int, default=None, help='maximum depth in tree ensemble.')

    # TREX settings
    parser.add_argument('--trex', action='store_true', default=False, help='use TREX.')
    parser.add_argument('--tree_kernel', type=str, default='leaf_output', help='type of encoding.')
    parser.add_argument('--true_label', action='store_true', default=False, help='train model on true labels.')
    parser.add_argument('--kernel_model', type=str, default='klr', help='kernel model to use.')

    # method settings
    parser.add_argument('--inf_k', type=int, default=None, help='number of leaves to use for leafinfluence.')
    parser.add_argument('--maple', action='store_true', default=False, help='whether to use MAPLE as a baseline.')
    parser.add_argument('--teknn', action='store_true', default=False, help='use KNN on top of TREX features.')
    parser.add_argument('--mmd', action='store_true', default=False, help='MMD-Critic prototypes.')
    parser.add_argument('--proto', action='store_true', default=False, help='Tree prototypes.')

    # plot settings
    parser.add_argument('--check_pct', type=float, default=0.3, help='max percentage of train instances to check.')
    parser.add_argument('--n_plot_points', type=int, default=10, help='number of points to plot.')

    # experiment settings
    parser.add_argument('--rs', type=int, default=1, help='random state.')
    parser.add_argument('--verbose', type=int, default=0, help='verbosity level.')

    args = parser.parse_args()
    main(args)


from wagtail.core.fields import StreamField, StreamValue
from graphene.types import Scalar

from graphene_django.converter import convert_django_field

class GenericStreamFieldType(Scalar):
    @staticmethod
    def serialize(stream_value: StreamValue):
        block = stream_value.stream_block
        return block.get_api_representation(stream_value)

@convert_django_field.register(StreamField)
def convert_stream_field(field: StreamField, registry=None):
    return GenericStreamFieldType(description=field.help_text,
                                  required=not field.null)


import numpy as np
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import argparse
import os
import os.path
import scipy
import scipy.stats

'''
Author: Fernando Yrazu.

Created: May 2014

Description:
 This program plots the ddG0, ddGdag and Phi-value comparison from simulation and experiment

Procedure:
   Execute the program with the following options:
   --prot (protein_name)
   --iter (iteration_number)
   The iteration number should start from zero (meaning, raw simulation without experimental adjustment)

Changelog:
May 2014 Created                                                                                June 2014 Added phi-value comparison                                                                                July 2014 Updated to new directory format, eliminated references to excel docs that were unnecessary
'''
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--prot', type=str, required=True, help='Input protein name')
    parser.add_argument('--iter', type=str, required=True, help='Input iteration number')
    args = parser.parse_args()

    return args

def plot_ddGs(protein, current_dir, iteration, select_temp):
    select_path_1 = current_dir+'/'+protein+'/Mut_'+iteration+'/'+select_temp+'_1/phi/'
    select_path_2 = current_dir+'/'+protein+'/Mut_'+iteration+'/'+select_temp+'_2/phi/'
    select_path_3 = current_dir+'/'+protein+'/Mut_'+iteration+'/'+select_temp+'_3/phi/'

    ddGdag_sim_1 = np.loadtxt(select_path_1+'Q_phi.dat', usecols=(4,))
    ddG0_sim_1 = np.loadtxt(select_path_1+'Q_phi.dat', usecols=(5,))
    ddGdag_sim_2 = np.loadtxt(select_path_2+'Q_phi.dat', usecols=(4,))
    ddG0_sim_2 = np.loadtxt(select_path_2+'Q_phi.dat', usecols=(5,))
    ddGdag_sim_3 = np.loadtxt(select_path_3+'Q_phi.dat', usecols=(4,))
    ddG0_sim_3 = np.loadtxt(select_path_3+'Q_phi.dat', usecols=(5,))

    ddGdag_sim = np.mean(np.column_stack((ddGdag_sim_1,ddGdag_sim_2,ddGdag_sim_3)), axis=1)
    ddG0_sim = np.mean(np.column_stack((ddG0_sim_1,ddG0_sim_2,ddG0_sim_3)), axis=1)
    err_ddGdag_sim = np.abs(np.amax(np.column_stack((ddGdag_sim_1,ddGdag_sim_2,ddGdag_sim_3)), axis=1)-
                            np.amin(np.column_stack((ddGdag_sim_1,ddGdag_sim_2,ddGdag_sim_3)), axis=1))
    err_ddG0_sim = np.abs(np.amax(np.column_stack((ddG0_sim_1,ddG0_sim_2,ddG0_sim_3)), axis=1)-
                          np.amin(np.column_stack((ddG0_sim_1,ddG0_sim_2,ddG0_sim_3)), axis=1))
    ddGdag_exp_raw = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(6,))
    err_ddGdag_exp_raw = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(7,))
    ddG0_exp_raw = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(4,))
    err_ddG0_exp_raw = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(5,))
    index_raw = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(1,))
    usable = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(8,), dtype=str)
    location = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(0,), dtype=str)
    exclude = np.loadtxt(current_dir+'/'+protein+'_calculated_ddG.dat', usecols=(11,), dtype=int)

    ddGdag_exp = []
    err_ddGdag_exp = []
    ddG0_exp = []
    err_ddG0_exp = []
    index =[]
    for i in range(len(usable)):
        if all([usable[i]=='True', location[i]=='core',exclude[i]==0]):
            ddGdag_exp.append(ddGdag_exp_raw[i])
            err_ddGdag_exp.append(err_ddGdag_exp_raw[i])
            ddG0_exp.append(ddG0_exp_raw[i])
            err_ddG0_exp.append(err_ddG0_exp_raw[i])
            index.append(index_raw[i])

    #Calculate Phi values
    phi_exp = []
    phi_sim = []
    phi_exp_rel = []
    phi_sim_rel = []
    phi_exp_raw = []
    phi_sim_raw = []
    # Approximate value in kT that makes a ddG0 usable for the Phi-value calculation
    ddG0_usable_minimum = 1.30

    for i in range(len(ddG0_exp)):
        exp = ddGdag_exp[i]/ddG0_exp[i]
        sim = ddGdag_sim[i]/ddG0_sim[i]
        # Approximate value in kT that
        if ddG0_exp[i]<ddG0_usable_minimum:
            phi_exp_rel.append(exp)
            phi_sim_rel.append(sim)
        else:
            phi_exp.append(exp)
            phi_sim.append(sim)
        #these are for making the bar plot later on
        phi_exp_raw.append(exp)
        phi_sim_raw.append(sim)

    #First plot
    #Finding the maxima and minima to size plot and draw y=x line
    c = np.min(ddGdag_sim)
    d = np.min(ddGdag_exp)
    e = np.max(ddGdag_sim)
    f = np.max(ddGdag_exp)
    g = np.min(ddG0_sim)
    h = np.min(ddG0_exp)
    m = np.max(ddG0_sim)
    n = np.max(ddG0_exp)
    upper = np.max([e,f,m,n]) + 1
    lower = np.min([c,d,g,h,0]) - 1
    fig = plt.figure()
    ax1 = fig.add_subplot(111)

#    ax1.plot(ddGdag_sim, ddGdag_exp, 'ro', label="ddGdag ")
    ax1.errorbar(ddGdag_sim, ddGdag_exp, fmt='ro', label="ddGdag ", xerr = err_ddGdag_sim, yerr= err_ddGdag_exp)
#    ax1.errorbar(ddGdag_sim, ddGdag_exp, fmt='ro', label="ddGdag ", xerr =0.08 , yerr=0.05)
#    ax1.plot(ddG0_sim, ddG0_exp, 'go', label="ddG0")
    ax1.errorbar(ddG0_sim, ddG0_exp, fmt='go', label="ddG0" , xerr = err_ddG0_sim, yerr=err_ddG0_exp)

    #Plot the tendency line of the entire sim vs exp ddGs (as in MatysiakClementi(20040
    ddG_sim = np.concatenate((ddGdag_sim,ddG0_sim))
    ddG_exp = np.concatenate((ddGdag_exp,ddG0_exp))
    correl , pval = scipy.stats.pearsonr(ddG_sim,ddG_exp)
    correlstr = "{:.6f}".format(correl)
    pvalstr = "{:.6f}".format(pval)
    fit = np.polyfit(ddG_sim,ddG_exp, 1)
    fit_fn = np.poly1d(fit)
    ax1.plot([lower,upper], fit_fn([lower,upper]), '-k')
    # Plot the 45 degree line
    ax1.plot([lower,upper],[lower,upper],'b-', lw=2)
    plt.title("Comparison of simulation and experimental ddG for "+protein+", iteration "+iteration+" \n Pearson correlation "+correlstr+" and p-value "+pvalstr)
    plt.xlim([lower,upper])
    plt.ylim([lower,upper])
    plt.xlabel('ddG from simulation')
    plt.ylabel('ddG from experiment')
    plt.legend(loc='lower right')
    plt.savefig(current_dir+'/metrics/ddG_comparison/'+protein+'_ddGs_'+iteration+'_'+select_temp+'.pdf')
    plt.clf()

    #Third plot: Phi values exp vs sim
    #Finding the maxima and minima to size plot and draw y=x line
    #Using latex for phi
#    plt.rc('text', usetex=True)

    c = np.min(phi_sim)
    d = np.min(phi_exp)
    if phi_exp_rel:
        e = np.max(phi_sim_rel)
        f = np.max(phi_exp_rel)
    else:
        e = 0.
        f = 0.

    upper = np.max([e,f]) + 1
    lower = np.min([c,d,0]) - 1

    plt.plot(phi_sim, phi_exp, 'ro')
    if phi_exp_rel:
        plt.plot(phi_sim_rel, phi_exp_rel, 'yo')

    #Plot the 45 degree line
    plt.plot([lower,upper],[lower,upper],'b-', lw=2)

    plt.title("Comparison of simulation and experimental $\Phi$ values for "+protein+", iteration "+iteration)
    plt.xlim([lower,upper])
    plt.ylim([lower,upper])
    plt.xlabel('$\Phi$ values from simulation')
    plt.ylabel('$\Phi$ values from experiment')
    plt.savefig(current_dir+'/metrics/ddG_comparison/'+protein+'_phi_values_'+iteration+'_'+select_temp+'.pdf')
    plt.clf()

    #Fourth plot: Phi values exp vs sim according to sequence
    #Finding the maxima and minima to size plot and draw y=x line
    c = np.min(phi_sim)
    d = np.min(phi_exp)
    if phi_exp_rel:
        e = np.max(phi_sim+phi_sim_rel)
        f = np.max(phi_exp+phi_exp_rel)
    else:
        e = np.max(phi_sim)
        f = np.max(phi_exp)

    upper = np.max([f]) + 1
    lower = np.min([d,0]) - 1
    upper_all = np.max([e,f]) + 1
    lower_all = np.min([c,d,0]) - 1

    colors_exp = []
    colors_sim = []
    for i in range(len(ddG0_exp)):
        if ddG0_exp[i]<ddG0_usable_minimum:
            colors_exp.append('y')
            colors_sim.append('y')
        else:
            colors_exp.append('r')
            colors_sim.append('b')

    f, ax = plt.subplots(3,sharex=True)

    ax[0].bar(index,phi_exp_raw, color=colors_exp)
    ax[1].bar(index,phi_sim_raw, color=colors_sim)
    ax[2].bar(index,phi_sim_raw, color=colors_sim)
    ax[0].set_title("Comparison of simulation and experimental $\Phi$ values for "+protein+", iteration "+iteration+'\n according to sequence')
    ax[0].axhline(y=0, xmin=0, xmax=1)
    ax[1].axhline(y=0, xmin=0, xmax=1)
    ax[2].axhline(y=0, xmin=0, xmax=1)
    ax[0].set_ylim([lower,upper])
    ax[1].set_ylim([lower,upper])
    ax[2].set_ylim([lower_all,upper_all])
    ax[0].set_ylabel('$\Phi$ values (exp)')
    ax[1].set_ylabel('$\Phi$ values (sim)')
    ax[2].set_ylabel('$\Phi$ values (sim)')
#    f.subplots_adjust(hspace=0)
#    plt.setp([a.get_xticklabels() for a in f.axes[:-1]], visible=False)
    plt.savefig(current_dir+'/metrics/ddG_comparison/'+protein+'_phi_value_sequence_'+iteration+'_'+select_temp+'.pdf')
    plt.clf()

    #Fifth plot: Eigenvalues of M matrix for decomposition
    select_path_mut = current_dir+'/'+protein+'/Mut_'+iteration+'/'+select_temp+'_agg/mut/singular_values_norm.dat'
    protein_dict={'r15':'ro', 'r16':'go', 'r17':'bo'}
    if protein in protein_dict:
        chosen_color = protein_dict[protein]
    else:
        chosen_color = 'ro'

    if os.path.exists(select_path_mut):
        ev = np.loadtxt(select_path_mut)
        plt.plot(ev, chosen_color)
        plt.axhline(y=0.5, xmin=0, xmax=1, color='m')
        plt.axhline(y=0.25, xmin=0, xmax=1, color='m')
        plt.axhline(y=0.1, xmin=0, xmax=1, color='m')
        plt.title("M-matrix normalized eigenvalues for "+protein)
        plt.xlabel('Ordinal #')
        plt.ylabel('Normalized eigenvalues')
        plt.savefig(current_dir+'/metrics/ddG_comparison/'+protein+'_'+select_temp+'_M_eigenvalues.pdf')

def main():
    protein = get_args().prot
    iteration = get_args().iter
    current_dir = os.getcwd()

    if os.path.isdir(current_dir+'/metrics')==False:
        os.mkdir('metrics')
    os.chdir(current_dir+'/metrics/')
    if os.path.isdir(current_dir+'/metrics/ddG_comparison')==False:
        os.mkdir('ddG_comparison')

    temps_file = open(current_dir+'/'+protein+'/Mut_'+iteration+'/Tf.txt').readlines()
    for temp in temps_file:
        select_temp = temp.split()[0]
        if not os.path.exists(current_dir+'/'+protein+'/Mut_'+iteration+'/'+select_temp+'_1'):
            temps_file_2 = open(current_dir+'/'+protein+'/Mut_'+iteration+'/ddG_Tf.txt').readlines()
            for temp2 in temps_file_2:
                select_temp = temp2.split()[0]
    plot_ddGs(protein, current_dir, iteration,select_temp)

if __name__=="__main__":
    main()


import warnings
import unittest
import zope.component.testing

from zope.interface import Interface, implements, alsoProvides
from zope.component import getMultiAdapter, provideUtility
from zope.configuration.config import ConfigurationExecutionError

from zope.publisher.browser import TestRequest

import grokcore.component.testing
import grokcore.view.interfaces

from zope.security.permission import Permission

from plone.z3cform.layout import FormWrapper

from plone.directives import form
from five import grok

import plone.directives.form.meta

class ILayer(Interface):
    pass

class IDummy(Interface):
    pass

class IDummy2(Interface):
    pass

class Dummy(object):
    implements(IDummy)

    def absolute_url(self):
        return "http://dummy"

class Request(TestRequest):

    # Zope 2 requests have this
    def __setitem__(self, name, value):
        self._environ[name] = value

# ignore warnings about unassociated templates, since the way we do the tests
# mean the association happens after the module is grokked
warnings.filterwarnings(
    "ignore",
    ".*unassociated template.*testformwithtemplate.*")

class TestFormDirectives(unittest.TestCase):

    def setUp(self):

        # On Zope 2.10, the default is True, on 2.12 it's False
        self.defaultWrap = plone.directives.form.meta.DEFAULT_WRAP
        plone.directives.form.meta.DEFAULT_WRAP = False

        grokcore.component.testing.grok('grokcore.component.meta')
        grokcore.component.testing.grok('grokcore.security.meta')
        grokcore.component.testing.grok('grokcore.view.meta.templates')
        grokcore.component.testing.grok('grokcore.view.meta.views')
        grokcore.component.testing.grok('grokcore.view.meta.skin')
        grokcore.component.testing.grok('five.grok.meta')
        grokcore.component.testing.grok('plone.directives.form.meta')

        grokcore.component.testing.grok('grokcore.view.templatereg')

        provideUtility(Permission('zope2.View'), name='zope2.View')
        provideUtility(
            Permission('cmf.ModifyPortalContent'),
            name='cmf.ModifyPortalContent')
        provideUtility(
            Permission('cmf.AddPortalContent'),
            name='cmf.AddPortalContent')

    def tearDown(self):
        zope.component.testing.tearDown()
        plone.directives.form.meta.DEFAULT_WRAP = self.defaultWrap

    def test_form_grokker_with_directives(self):

        class TestForm(form.Form):
            grok.context(IDummy)
            grok.name('my-test-form')
            grok.layer(ILayer)
            grok.require('zope2.View')

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()
        alsoProvides(request, ILayer)

        view = getMultiAdapter((context, request), name="my-test-form")

        self.failUnless(issubclass(view.form, TestForm))

    def test_grokker_with_defaults(self):

        class TestForm(form.Form):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))

    def test_schema_form(self):

        class TestForm(form.SchemaForm):
            grok.context(IDummy)
            schema = IDummy2

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.assertEquals(IDummy2, view.form.schema)

    def test_schema_form_implicit_schema(self):

        class TestForm(form.SchemaForm):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.assertEquals(IDummy, view.form.schema)

    def test_add_form(self):

        class TestForm(form.AddForm):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))

        self.assertEquals("http://dummy", view.form_instance    .nextURL())
        view.form_instance.immediate_view = "http://other_view"
        self.assertEquals("http://other_view", view.form_instance.nextURL())

    def test_schema_add_form(self):

        class TestForm(form.SchemaAddForm):
            grok.context(IDummy)
            schema = IDummy2

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.assertEquals(IDummy2, view.form.schema)

    # Note: No implicit schema-from-context here, since context of add form
    #  is container.

    def test_edit_form(self):

        class TestForm(form.EditForm):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))

    def test_schema_edit_form(self):

        class TestForm(form.SchemaEditForm):
            grok.context(IDummy)
            schema = IDummy2

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.assertEquals(IDummy2, view.form.schema)

    def test_schema_edit_form_implicit_schema(self):

        class TestForm(form.SchemaEditForm):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.assertEquals(IDummy, view.form.schema)

    def test_display_form_with_schema(self):

        # Note: We're not testing view registration, since that is done
        # by the standard five.grok View grokker

        class TestForm(form.DisplayForm):
            grok.context(IDummy)
            schema = IDummy2

        grokcore.component.testing.grok_component('TestForm', TestForm)

        self.assertEquals(IDummy2, TestForm.schema)

    def test_display_form_implicit_schema(self):

        # Note: We're not testing view registration, since that is done
        # by the standard five.grok View grokker

        class TestForm(form.DisplayForm):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        self.assertEquals(IDummy, TestForm.schema)

    def test_wrap(self):

        class TestForm(form.Form):
            grok.context(IDummy)
            form.wrap()

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.failUnless(isinstance(view, FormWrapper))

    def test_wrap_false(self):

        class TestForm(form.Form):
            grok.context(IDummy)
            form.wrap(False)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.failUnless(isinstance(view, TestForm))

    def test_nowrap(self):

        class TestForm(form.Form):
            grok.context(IDummy)
            form.wrap(False)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.failUnless(isinstance(view, TestForm))

    def test_wrap_default(self):

        # Simulate Zope 2.10 default
        plone.directives.form.meta.DEFAULT_WRAP = True

        class TestForm(form.Form):
            grok.context(IDummy)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.failUnless(isinstance(view, FormWrapper))

    def test_wrap_default_wrap_false(self):
        # Simulate Zope 2.10 default
        plone.directives.form.meta.DEFAULT_WRAP = True

        class TestForm(form.Form):
            grok.context(IDummy)
            form.wrap(False)

        grokcore.component.testing.grok_component('TestForm', TestForm)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testform")

        self.failUnless(issubclass(view.form, TestForm))
        self.failUnless(isinstance(view, TestForm))

    def test_template(self):

        class TestFormWithTemplate(form.Form):
            grok.context(IDummy)

        grokcore.component.testing.grok(__name__)

        grokcore.component.testing.grok_component(
            'TestFormWithTemplate',
            TestFormWithTemplate)

        context = Dummy()
        request = Request()

        view = getMultiAdapter((context, request), name="testformwithtemplate")

        self.failUnless(issubclass(view.form, TestFormWithTemplate))
        self.failUnless(
            grokcore.view.interfaces.ITemplate.providedBy(
                view.template))

    def test_template_and_render(self):

        class TestFormWithTemplate(form.Form):
            grok.context(IDummy)

            def render(self):
                return u"My custom renderer"

        grokcore.component.testing.grok(__name__)

        self.assertRaises(ConfigurationExecutionError,
                          grokcore.component.testing.grok_component,
                          'TestFormWithTemplate', TestFormWithTemplate)

def test_suite():
    return unittest.defaultTestLoader.loadTestsFromName(__name__)


######################################
#
# https://gist.github.com/KhyatiMahendru/6bf974c16bb9ce2bb8e8deedac03b01d#file-elbow_method-py
#
# por: KhyatiMahendru
# acesso em: 16/05/2020
#
######################################

from sklearn.cluster import KMeans

# function returns WSS score for k values from 1 to kmax
def calculate_WSS(points, kmax):
  sse = []
  for k in range(1, kmax+1):
    kmeans = KMeans(n_clusters = k).fit(points)
    centroids = kmeans.cluster_centers_
    pred_clusters = kmeans.predict(points)
    curr_sse = 0

    # calculate square of Euclidean distance of each point from its cluster center and add to current WSS
    for i in range(len(points)):
      curr_center = centroids[pred_clusters[i]]
      curr_sse += (points[i, 0] - curr_center[0]) ** 2 + (points[i, 1] - curr_center[1]) ** 2

    sse.append(curr_sse)
  return sse


import pandas as pd
import rdkit
from rdkit import Chem
from rdkit.Chem import AllChem
import argparse

#Define the src and target for torchtext to process
def run_smiles_generator(smiles):
    fp = []
    fpz = ""
    for i in range(256):
        fpz += '0'

    for i in range(len(smiles)):
        m = Chem.MolFromSmiles(smiles[i])
        try:
            fp1 = AllChem.GetMorganFingerprintAsBitVect(m, 2, nBits=256)
            fp1 = fp1.ToBitString()
            fp.append(fp1)
        except:
            fp.append(fpz)

    data = []
    for i in range(len(smiles)):
        t = []
        for j in range(len(fp[i])):
            t.append(fp[i][j])
        data.append(t)

    f1 = pd.DataFrame(data)
    st = 'MFP_'
    for i in range(256):
        col = st + str(i)
        f1=f1.rename(columns = {i:col})

    return(f1)

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", help = "input filename")
    parser.add_argument("--output", help = "output filename")
    args = parser.parse_args()

    print('Inputfile:',args.input)
    print('Outputfile:',args.output)

    #Read input file containing uniprot_id, protein sequence, inchikey, smiles, pchembl value
    input_df = pd.read_csv("../data/"+args.input,header='infer')
    all_sequences = input_df['canonical_smiles'].values.tolist()

    #Pass the file to Autoencoder and get latent space dataframe
    ls_df = run_smiles_generator(all_sequences)

    #Write output
    ls_df.to_csv("../data/"+args.output,index=False)
    print("Morgan fingerprints generation done")


import os
import PIL.Image as Image

# array = []
# for i in range(1, 10):
#     array.append(10 - i)
# array.sort(lambda(lhs, rhs): return lhs-rhs)
# print(array)

allSize = 0

def PrintImgSize(imgPath):
    if not os.path.exists(imgPath):
        print("not found => " + imgPath)
        return
    global allSize
    imgData = Image.open(imgPath)
    # print(imgData.size)
    print(imgPath, imgData.size[0], imgData.size[1])
    allSize = allSize + imgData.size[0] * imgData.size[1]
    imgData.close()

def Load(resName, sex, direction, action):
    respath = "E:\\lycq\\resource\\assets\\movie\\body\\" + ("body" + resName + "_" + sex + "_" + direction + action) + ".png"
    PrintImgSize(respath)

def Load2(resName):
    for sex in range(0, 2):
        for dir in range(0, 5):
            Load(resName, str(sex), str(dir), "a")
            Load(resName, str(sex), str(dir), "c")
            Load(resName, str(sex), str(dir), "r")
            Load(resName, str(sex), str(dir), "s")

# Load2("000")
# Load2("001")
# Load2("002")
# for i in range(1, 13):
#     Load2(str(100 + i))

# for i in range(1, 13):
#     Load2(str(200 + i))

# for i in range(1, 13):
#     Load2(str(300 + i))

# def LoadHero(type, resName, direction, action):
#     respath = "E:\\lycq\\resource\\assets\\movie\\" + type + "\\" + (type + resName + "_" + direction + action) + ".png"
#     PrintImgSize(respath)

# def LoadHero2(type, resName):
#         for dir in range(0, 5):
#             LoadHero(type, resName, str(dir), "a")
#             LoadHero(type, resName, str(dir), "c")
#             LoadHero(type, resName, str(dir), "r")
#             LoadHero(type, resName, str(dir), "s")

# for i in range(1, 10):
#     LoadHero2("hero", str(1000 + i))

# for i in range(1, 10):
#     LoadHero2("monster", str(10000 + i))

def GetAllImgPath(fileDir):
    imgList = []
    # for fileName in os.listdir(fileDir):
    for fileDir, sub, fileList in os.walk(fileDir):
        for fileName in fileList:
            if fileName.endswith("png") or fileName.endswith(".jpg"):
                imgList.append(os.path.join(fileDir, fileName))

    # print(os.walk(fileDir))
        # if fileName.endswith("png"):
        #     imgList.append(os.path.join(fileDir, fileName))
    return imgList

def PrintSize(dir):
    for path in GetAllImgPath(dir):
        PrintImgSize(path)

PrintSize("E:\\lycq\\resource\\assets\\movie\\body")
# PrintSize("E:\\lycq\\resource\\assets\\movie\\weapon")
# PrintSize("E:\\lycq\\resource\\assets\\movie\\wing")
# PrintSize("E:\\lycq\\resource\\temp_out_movie\\wing")
# PrintSize("F:\\game\\release\\~213\\resource\\assets\\atlas2_ui")
# PrintSize("E:\\lycq\\resource\\temp_out_movie\\body")

print(allSize * 1)

# 29 541 516


from django.http import Http404, HttpResponse, JsonResponse
from django.shortcuts import render
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import csrf_exempt
from django.views.generic import View

from .forms import UploadCameraImageForm
from .models import Camera, CameraImage

class Homepage(View):
    template_name = "pages/home.html"

    def get(self, request):
        try:
            cameras = Camera.objects.exclude(slug='ccf')
        except Camera.DoesNotExist:
            raise Http404("No cameras available.")

        return render(request, self.template_name, {
            'cameras': cameras,
            'camera': cameras[0],
        })

class ViewCamera(View):
    template_name = "pages/camera.html"

    def get(self, request, camera_slug):
        try:
            camera = Camera.objects.get(slug=camera_slug)
        except Camera.DoesNotExist:
            raise Http404("Camera does not exist")

        return render(request, self.template_name, {
            'camera': camera,
            'cameras': [camera, ],
            'single': True,
        })

@method_decorator(csrf_exempt, name='dispatch')
class UploadCameraImage(View):
    def post(self, request, secret):
        form = UploadCameraImageForm(request.POST, request.FILES)

        if form.is_valid():
            camera = None
            try:
                camera = Camera.objects.get(secret=secret)
            except Camera.DoesNotExist:
                pass

            if not camera:
                try:
                    camera = Camera.objects.get(ftp_username=secret)
                except Camera.DoesNotExist:
                    pass

            if not camera:
                raise Http404("Camera does not exist")

            image = CameraImage.objects.create(
                camera=camera,
                image=request.FILES['image'],
            )

            return HttpResponse('ok: #{}'.format(image.id))
        else:
            print(form.errors)
            raise Http404("Invalid image")

class GetCameraImage(View):
    def get(self, request, camera_id, width, height):
        try:
            camera = Camera.objects.get(id=camera_id)
        except Camera.DoesNotExist:
            return JsonResponse({'error': "Camera does not exist"})

        if not camera.current_image:
            return JsonResponse({'error': "No camera image (image not set)"})

        return JsonResponse(camera.current_image.get_sized_image(width, height))


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

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('turnos', '0007_auto_20150222_2109'),
    ]

    operations = [
        migrations.AlterField(
            model_name='eventid',
            name='event',
            field=models.OneToOneField(related_name='event_number', primary_key=True, serialize=False, to='turnos.Event'),
            preserve_default=True,
        ),
    ]


import os
import sys
import json
from glob import glob
import ast
import csv

import mylib

pathInput = "data/input/"
#inputTsv = pathInput + "citationNet_v3_" + "_".join(mylib.sectors).replace("/","") + ".json"
inputTsv = pathInput + "_".join(mylib.sectors).replace("/","") + "_withNames_scopusId_editedAddedMissingScopusId.tsv"

pathOutput = "data/output/"
pathAbstracts = pathOutput + "abstracts/"+ ("_".join(mylib.sectors)).replace("/", "") + "/"
fileCitingNet = pathOutput + "citingNet_" + ("_".join(mylib.sectors)).replace("/", "") + ".json"
codePrefix = "2-s2.0-"
fileDoiEidMap = pathOutput + "doiEidMap_" + ("_".join(mylib.sectors)).replace("/", "") + ".json"
eidToAuthorsFile = pathOutput + "eidToAuthors_" + ("_".join(mylib.sectors)).replace("/", "") + ".json"

outputTsv = pathOutput + "adjacencyMatrix_" + ("_".join(mylib.sectors)).replace("/", "") + ".tsv"

class AutoVivification(dict):
    """Implementation of perl's autovivification feature."""
    def __getitem__(self, item):
        try:
            return dict.__getitem__(self, item)
        except KeyError:
            value = self[item] = type(self)()
            return value

def computeCitedNetwork(path):
	documentSet = set()
	contents = glob(path + '*.json')
	contents.sort()
	# get set of eids from the input list of dois
	for filename_withPath in contents:
		with open(filename_withPath) as json_file:
			data = json.load(json_file)
			try:
				citing = data['abstracts-retrieval-response']['coredata']['eid']
				documentSet.add(citing)
			except:
				print ('*** EID ERROR: ' + filename_withPath + " ***")

	# create the dictionary of cited papers (with empty lists of citing papers)
	citingCitedDict = dict()
	for citing in documentSet:
		citingCitedDict[citing] = []

	for filename_withPath in contents:
		with open(filename_withPath) as json_file:
			data = json.load(json_file)
			try:
				citing = data['abstracts-retrieval-response']['coredata']['eid']
				biblio = data['abstracts-retrieval-response']['item']['bibrecord']['tail']['bibliography']
				numRefs = int(biblio['@refcount'])
				for ref in biblio['reference']: #range(numRefs):
					temp = ref['ref-info']['refd-itemidlist']['itemid']
					#print (filename_withPath)
					if type(temp) is not list:
						idtype = temp['@idtype']
						code = temp['$']
						cited = codePrefix + code
						if cited in documentSet:
							citingCitedDict[citing].append(cited)
					else:
						for el in temp:
							idtype = el['@idtype']
							code = el['$']
							cited = codePrefix + code
							if idtype == 'SGR':
								if cited in documentSet:
									citingCitedDict[citing].append(cited)
			except:
				print ('*** BIBLIOGRAPHY ERROR: ' + filename_withPath + " ***")

	return citingCitedDict

def getEidToAuthors(tsv, doiEidMap):
	eidToAuthors = dict()
	with open(tsv, newline='') as csvfile:
		spamreader = csv.DictReader(csvfile, delimiter='\t')
		for row in spamreader:
			authorId = row['SCOPUS ID']
			citingDois = ast.literal_eval(row['DOIS ESISTENTI'])
			for citingDoi in citingDois:
				try:
					citingEid = doiEidMap["doi-to-eid"][citingDoi]
					if citingEid in eidToAuthors:
						temp = eidToAuthors[citingEid]
						if authorId not in temp:
							temp.append(authorId)
							eidToAuthors[citingEid] = temp
					else:
						eidToAuthors[citingEid] = [authorId]
				except:
					continue
	return eidToAuthors

def getAuthorIdsList(tsv):
	authorIdsSet = set()
	with open(tsv, newline='') as csvfile:
		spamreader = csv.DictReader(csvfile, delimiter='\t')
		for row in spamreader:
			authorId = row['SCOPUS ID']
			authorIdsSet.add(authorId)
	authorIdsList = list(authorIdsSet)
	authorIdsList.sort()
	return authorIdsList

def getAdjacencyMatrix(tsv, authorIdsList, doiEidMap, citingCitedDict, eidToAuthors):

	# initialize empty map (dict of dict)
	matrix = AutoVivification()
	for authorId1 in authorIdsList:
		for authorId2 in authorIdsList:
			matrix[authorId1][authorId2] = 0

	with open(tsv, newline='') as csvfile:
		authorsAnalyzed = set()
		spamreader = csv.DictReader(csvfile, delimiter='\t')
		for row in spamreader:
			citingAuthor = row['SCOPUS ID']
			if citingAuthor in authorsAnalyzed:
				continue
			authorsAnalyzed.add(citingAuthor)
			citingDois = ast.literal_eval(row['DOIS ESISTENTI'])
			for citingDoi in citingDois:
				try:
					citingEid = doiEidMap["doi-to-eid"][citingDoi]
					for citedEid in citingCitedDict[citingEid]:
						citedAuthors = eidToAuthors[citedEid]
						for citedAuthor in citedAuthors:
							temp = matrix[citingAuthor][citedAuthor]
							#print (temp)
							temp += 1
							matrix[citingAuthor][citedAuthor] = temp
				except:
					continue
	return matrix

#Save Adjacency Matrix to TSV file
def saveMatrixToFile(tsv, matrix, authorList):
	res = "\tCV-" + "\tCV-".join(authorList) + "\n"
	for authorIdCiting in authorList:
		line = list()
		for authorIdCited in authorList:
			line.append(str(matrix[authorIdCiting][authorIdCited]))
		res += "CV-" + authorIdCiting + "\t" + "\t".join(line) + "\n"
	text_file = open(tsv, "w")
	text_file.write(res)
	text_file.close()

# load doiEid map from file
with open(fileDoiEidMap) as json_file:
	doiEidMap = json.load(json_file)

eidToAuthors = getEidToAuthors(inputTsv, doiEidMap)

with open(eidToAuthorsFile, 'w') as fp:
	json.dump(eidToAuthors, fp, indent=4)

# compute citing network (eid-citing -> list-eids-cited) using abstracts (i.e. json files)
citingCitedDict = computeCitedNetwork(pathAbstracts)

# save the citing network (eid-citing -> list-eids-cited) to file
with open(fileCitingNet, 'w') as fp:
	json.dump(citingCitedDict, fp, indent=4)

# load the citing network (eid-citing -> list-eids-cited) from file
#with open(fileCitingNet) as json_file:
#	citingCitedDict = json.load(json_file)

authorIdsList = getAuthorIdsList(inputTsv)

adjacencyMatrix = getAdjacencyMatrix(inputTsv, authorIdsList, doiEidMap, citingCitedDict, eidToAuthors)

saveMatrixToFile(outputTsv, adjacencyMatrix, authorIdsList)

'''
with open(inputTsv) as json_file:
	data = json.load(json_file)
	# get author ids
	authorIds = list()
	for authorId in data:
		authorIds.append("cv-" +authorId)

	# initialize empty map (dict of dict)
	matrix = AutoVivification()
	for authorId1 in authorIds:
		for authorId2 in authorIds:
			matrix[authorId1][authorId2] = 0
			print (matrix[authorId1][authorId2])

	for citing in citingCitedDict:
		for cited in citingCitedDict[citing]:
			temp = matrix[citing][cited]
			print (temp)
			matrix[citing][cited] = temp + 1

	res = "\t" + "\t".join(authorIds) + "\n"
	for authorIdCiting in authorIds:
		res += authorIdauthorIdCiting
		for authorIdCited in authorIds:
			res += matrix[authorIdCiting][authorIdCited]
		res += "\n"

	print (res)

	citingMap = dict()
	for authorCited in data:
		for idCv in data[authorCited]:
			for pub in data[authorCited][idCv]["pubs"]:
				cited = pub["cited"]["doi"]
				citing = pub["citing"]["doi"]
				authorsCiting = pub["citing"]["authors"]

				for authorCitingDict in authorsCiting:
					authorCiting = authorCitingDict["authorId"]
					temp = dict()
					if authorCiting in citingMap:
						temp = citingMap[authorCited]
'''

'''
		c1	c2	...	cn
	c1
	c2
	...
	cn
'''