import pybullet as p
# from time import sleep
import pybullet_data
import numpy as np
import matplotlib.pyplot as plt
import time
import sys
import make_body_plot as m

physicsClient = p.connect(p.DIRECT)
p.resetSimulation(p.RESET_USE_DEFORMABLE_WORLD)

p.setAdditionalSearchPath(pybullet_data.getDataPath())
planeId = p.loadURDF("plane.urdf", [0,0,0],globalScaling=100)
p.setGravity(0,0,-10)
sF = 4
arm_height = 5./sF
arm_width = 1./sF
arm_length = 4./sF
body_height = 1./sF
body_width = 3./sF
body_length = 4./sF

# create bodies
mod1 = p.loadURDF("cube_small.urdf",[body_length/2+arm_width+arm_length+0.25,body_width/2-arm_width/2,arm_height/2],globalScaling=5./sF)
p.changeVisualShape(mod1,-1,rgbaColor=[1,0,0,1])
mod2 = p.loadURDF("cube_small.urdf",[body_length/2+arm_width+arm_length+0.25,-body_width/2+arm_width/2,arm_height/2],globalScaling=5./sF)
p.changeVisualShape(mod2,-1,rgbaColor=[0,1,0,1])
mod3 = p.loadURDF("cube_small.urdf",[-arm_length-body_length/2-arm_width/2-0.75,body_width/2-arm_width/2,arm_height/2],globalScaling=5./sF)
p.changeVisualShape(mod3,-1,rgbaColor=[1,0,1,1])
mod4 = p.loadURDF("cube_small.urdf",[-arm_length-body_length/2-arm_width/2-0.75,-body_width/2+arm_width/2,arm_height/2],globalScaling=5./sF)
p.changeVisualShape(mod4,-1,rgbaColor=[0,0,1,1])

corrFac = 0.05/2
bodyId = p.loadURDF("body.urdf", [0,0,arm_height/2],globalScaling=1./sF)
armId1 = p.loadSoftBody("arm_hori.obj", basePosition = [body_length/2+arm_width/2+corrFac,body_width/2-arm_width/2,arm_height/2], scale = 0.1/sF, mass = 1./sF, useNeoHookean = 1, useBendingSprings=1, useMassSpring=1, springElasticStiffness=121, springDampingStiffness=1, springDampingAllDirections = 1, useSelfCollision = 1, frictionCoeff = 0, useFaceContact=1)
armId2 = p.loadSoftBody("arm_hori.obj", basePosition = [body_length/2+arm_width/2+corrFac,-body_width/2+arm_width/2,arm_height/2], scale = 0.1/sF, mass = 1./sF, useNeoHookean = 1, useBendingSprings=1, useMassSpring=1, springElasticStiffness=121, springDampingStiffness=1, springDampingAllDirections = 1, useSelfCollision = 1, frictionCoeff = 0, useFaceContact=1)
armId3 = p.loadSoftBody("arm_hori.obj", basePosition = [-arm_length-body_length/2-arm_width/2-corrFac,body_width/2-arm_width/2,arm_height/2], scale = 0.1/sF, mass = 1./sF, useNeoHookean = 1, useBendingSprings=1, useMassSpring=1, springElasticStiffness=121, springDampingStiffness=1, springDampingAllDirections = 1, useSelfCollision = 1, frictionCoeff = 0, useFaceContact=1)
armId4 = p.loadSoftBody("arm_hori.obj", basePosition = [-arm_length-body_length/2-arm_width/2-corrFac,-body_width/2+arm_width/2,arm_height/2], scale = 0.1/sF, mass = 1./sF, useNeoHookean = 1, useBendingSprings=1, useMassSpring=1, springElasticStiffness=121, springDampingStiffness=1, springDampingAllDirections = 1, useSelfCollision = 1, frictionCoeff = 0, useFaceContact=1)

p.changeVisualShape(bodyId,-1,rgbaColor=[160/255,69/255,19/255,1])
p.changeVisualShape(armId1,-1,rgbaColor=[0,0,0,0.8])
p.changeVisualShape(armId2,-1,rgbaColor=[0,0,0,0.8])
p.changeVisualShape(armId3,-1,rgbaColor=[0,0,0,0.8])
p.changeVisualShape(armId4,-1,rgbaColor=[0,0,0,0.8])

# attach bodies
p.createSoftBodyAnchor(armId1,7,bodyId,-1)
p.createSoftBodyAnchor(armId1,8,bodyId,-1)
p.createSoftBodyAnchor(armId1,16,bodyId,-1)
p.createSoftBodyAnchor(armId1,17,bodyId,-1)

p.createSoftBodyAnchor(armId2,7,bodyId,-1)
p.createSoftBodyAnchor(armId2,8,bodyId,-1)
p.createSoftBodyAnchor(armId2,16,bodyId,-1)
p.createSoftBodyAnchor(armId2,17,bodyId,-1)

p.createSoftBodyAnchor(armId3,0,bodyId,-1)
p.createSoftBodyAnchor(armId3,1,bodyId,-1)
p.createSoftBodyAnchor(armId3,2,bodyId,-1)
p.createSoftBodyAnchor(armId3,3,bodyId,-1)

p.createSoftBodyAnchor(armId4,0,bodyId,-1)
p.createSoftBodyAnchor(armId4,1,bodyId,-1)
p.createSoftBodyAnchor(armId4,2,bodyId,-1)
p.createSoftBodyAnchor(armId4,3,bodyId,-1)

p.createSoftBodyAnchor(armId1,0,mod1,-1)
p.createSoftBodyAnchor(armId1,1,mod1,-1)
p.createSoftBodyAnchor(armId1,2,mod1,-1)
p.createSoftBodyAnchor(armId1,3,mod1,-1)

p.createSoftBodyAnchor(armId2,0,mod2,-1)
p.createSoftBodyAnchor(armId2,1,mod2,-1)
p.createSoftBodyAnchor(armId2,2,mod2,-1)
p.createSoftBodyAnchor(armId2,3,mod2,-1)

p.createSoftBodyAnchor(armId3,7,mod3,-1)
p.createSoftBodyAnchor(armId3,8,mod3,-1)
p.createSoftBodyAnchor(armId3,16,mod3,-1)
p.createSoftBodyAnchor(armId3,17,mod3,-1)

p.createSoftBodyAnchor(armId4,7,mod4,-1)
p.createSoftBodyAnchor(armId4,8,mod4,-1)
p.createSoftBodyAnchor(armId4,16,mod4,-1)
p.createSoftBodyAnchor(armId4,17,mod4,-1)

numLinks = p.getMeshData(armId1)
pos,ort = p.getBasePositionAndOrientation(armId1)

# run simulation
useRealTimeSimulation = 0
if (useRealTimeSimulation):
  p.setRealTimeSimulation(0)

posArr1 = []
posArr2 = []
posArr3 = []
posArr4 = []
posArr5 = []

velArr1 = []
velArr2 = []
velArr3 = []
velArr4 = []

mods = [mod1,mod2,mod3,mod4]

# Control related variables
k_p = 3 # Proportional constant
k_i = 0.5 # Integral constant
k_d = 0.01 # Differential constant
# goal_vel_x = [0,0,0,0] # Velocity goal of each wheel
goal_vel_x = np.array([25,25,25,25])*4/25 # Velocity goal of each wheel
goal_pos_y = np.array([1.,-1.,1.,-1.])/sF # Position goal of each wheel
goal_vel_y = np.array([1,1,1,1])*0
# goal_pos_y = [0,0,0,0] # Velocity goal of each wheel
lastControlTime = 0
# target_vel = (50,10,0)
# turns = 1
control_x_prev = np.array([0,0,0,0])
# control_y_prev = [0,0,0,0]
control_y_prev = np.array([1,-1,1,-1])*0/sF
error_x_prev = [0,0,0,0]
error_y_prev = [0,0,0,0]
time_prev = time.time()
argv = sys.argv

# For velocity plot, goal vel is 25, kp is 10, no kd ki, both x y same k no  multiplier, y vel control,run for 10s, no need different force 12 vs 34
# For bad disturbance, noise use 3std, kp 10, pos control y, force divide by 5. x P y P
# For good disturbance, noise 3std, kp 10 kd 0.01, force divide by 3, x P y PD

def update_goal(t,op=None):
  global goal_vel_x,start_time,goal_pos_y
  if op == None:
    return
  elif op == 's':#sinusoidal with const x
    speed = 15*np.cos(2*np.pi*(t-start_time)/5)*1
    goal_vel_x = [speed,speed,speed,speed]
    goal_pos_y += np.array([1,1,1,1])*0.02
    return
  elif op == 'c_up':
    elapsed = t-start_time
    start = 3
    diff = 5
    init_vel = 10
    #start curve at t = 3 sec
    if elapsed < start:
      goal_vel_x = np.array([10,10,10,10])
    if elapsed > start and elapsed < start+diff:
      speed = init_vel*np.cos((t-start_time-start)*np.pi/10)
      goal_vel_x = [speed,speed,speed,speed]
      pos = -20*np.sin((t-start_time-start)*np.pi/10)
      goal_pos_y = np.array([1,-1,1,-1])+np.array([pos,pos,pos,pos])

    elif elapsed > start+diff:
      goal_pos_y += np.array([1,1,1,1])*0.01
      goal_vel_x = np.zeros(4)
    print(goal_pos_y)

start_time = time.time()
store_data_time = time.time()
for i in range(4):
  print(p.getBasePositionAndOrientation(mods[i]))
while time.time()-start_time < 0:
  p.stepSimulation()

start_time = time.time()
save=False

vel_x = np.zeros(4)

if len(sys.argv) > 3:
  k_p = float(sys.argv[1])
  k_i = float(sys.argv[2])
  k_d = float(sys.argv[3])
  save=False

endtime = 8
numits = 0
while numits < 240*endtime:
  numits += 1
  print(numits)
  for i in range(4):
    p.applyExternalForce(mods[i],-1,np.array([10,0,0]),[0,0,0],p.LINK_FRAME)
    vel_x[i] = p.getBaseVelocity(mods[i])[0][0]

  # Store
  pos1,ort = p.getBasePositionAndOrientation(mod1)
  pos2,ort = p.getBasePositionAndOrientation(mod2)
  pos3,ort = p.getBasePositionAndOrientation(mod3)
  pos4,ort = p.getBasePositionAndOrientation(mod4)
  pos5,ort = p.getBasePositionAndOrientation(bodyId)
  posArr1.append(pos1)
  posArr2.append(pos2)
  posArr3.append(pos3)
  posArr4.append(pos4)
  posArr5.append(pos5)

  velArr1.append(vel_x[0])
  velArr2.append(vel_x[1])
  velArr3.append(vel_x[2])
  velArr4.append(vel_x[3])
  p.stepSimulation()

posArr1 = np.array(posArr1)
posArr2 = np.array(posArr2)
posArr3 = np.array(posArr3)
posArr4 = np.array(posArr4)
posArr5 = np.array(posArr5)

    # Master array for position
masterArr = np.zeros([posArr1.shape[0],posArr1.shape[1],5])
masterArr[:,:,0] = posArr1
masterArr[:,:,1] = posArr2
masterArr[:,:,2] = posArr3
masterArr[:,:,3] = posArr4
masterArr[:,:,4] = posArr5

    # Master array for velocity in x direction
masterArr2 = np.zeros([len(velArr1),4])
masterArr2[:,0] = velArr1
masterArr2[:,1] = velArr2
masterArr2[:,2] = velArr3
masterArr2[:,3] = velArr4

'''
colors = ['red','green','purple','blue','black']
label = ['limb 1','limb 2','limb 3','limb 4']
fig, (ax1, ax2) = plt.subplots(2)
for i in range(4):
  test = ax1.plot(masterArr[:,0,i],masterArr[:,1,i],color=colors[i], label=label[i])
  ax1.scatter(masterArr[0,0,i],masterArr[0,1,i],color='black')
ax1.legend()
for i in range(4):
  ax2.plot(np.linspace(0, 10, num=len(velArr1)), masterArr2[:,i],color=colors[i], label=label[i])
ax1.set(xlabel='x position (m)', ylabel='y position (m)')
ax2.set(xlabel='time (s)', ylabel='x velocity (m/s)')
ax2.legend()
'''
np.save("limbPos.npy",masterArr)
np.save("limbVels.npy",masterArr2)
fig,ax1,ax2 = m.make_graded_limb_plot(masterArr,masterArr2,0,endtime,"Limb Position: Constant Force","Limb X Velocities vs Time: Constant Force")
fig2,ax3 = m.make_plot(posArr5,0,endtime)
if save:
  title=argv[1]+'kp_'+argv[2]+'ki_'+argv[3]+'kd.png'
  plt.savefig(title)
  raise(EOFError)
plt.show()
# raise(EOFError)


from flask import Flask
from flask_sqlalchemy import SQLAlchemy
from flask_bcrypt import Bcrypt
from flask_login import LoginManager

# Initialize SQLAlchemy and Bcrypt extensions (not yet associated with any Flask app).
db = SQLAlchemy()
bcrypt = Bcrypt()

def create_app(config_name=None):
    """
    Create an instance of the Flask application.
    """
    app = Flask(__name__)

    # A secret key for signing session cookies
    app.config["SECRET_KEY"] = "93220d9b340cf9a6c39bac99cce7daf220167498f91fa"

    # Application configuration
    app.config["SQLALCHEMY_DATABASE_URI"] = "mysql+mysqldb://24_webapp_031:dGlvdf65@mysql.lab.it.uc3m.es/24_webapp_031c"
    # Uncomment the following line to use MySQL
    # app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///recipes.db'
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = True

    db.init_app(app)

    # Login manager.
    login_manager = LoginManager()
    login_manager.login_view = 'auth.login'
    login_manager.init_app(app)

    # User loader for login management
    from . import model
    @login_manager.user_loader
    def load_user(user_id):
        return db.session.get(model.User, int(user_id))

    # Register blueprints
    from . import main, auth, explore
    app.register_blueprint(main.bp)
    app.register_blueprint(auth.bp)
    app.register_blueprint(explore.bp)

    return app


import os
import pickle
import numpy as np
import collections

def traverse_dir(
        root_dir,
        extension=('mid', 'MID'),
        amount=None,
        str_=None,
        is_pure=False,
        verbose=False,
        is_sort=False,
        is_ext=True):
    if verbose:
        print('[*] Scanning...')
    file_list = []
    cnt = 0
    for root, _, files in os.walk(root_dir):
        for file in files:
            if file.endswith(extension):
                if (amount is not None) and (cnt == amount):
                    break
                if str_ is not None:
                    if str_ not in file:
                        continue
                mix_path = os.path.join(root, file)
                pure_path = mix_path[len(root_dir)+1:] if is_pure else mix_path
                if not is_ext:
                    ext = pure_path.split('.')[-1]
                    pure_path = pure_path[:-(len(ext)+1)]
                if verbose:
                    print(pure_path)
                file_list.append(pure_path)
                cnt += 1
    if verbose:
        print('Total: %d files' % len(file_list))
        print('Done!!!')
    if is_sort:
        file_list.sort()
    return file_list

if __name__ == '__main__':
    # paths
    path_root = 'ailab17k_from-scratch_cp'
    path_indir = os.path.join(path_root, 'events')
    path_outdir =  os.path.join(path_root, 'words')
    path_dictionary = os.path.join(path_root, 'dictionary.pkl')
    os.makedirs(path_outdir, exist_ok=True)

    # list files
    eventfiles = traverse_dir(
        path_indir,
        is_pure=True,
        is_sort=True,
        extension=('pkl'))
    n_files = len(eventfiles)
    print('num fiels:', n_files)

    # --- build dictionary --- #
    # all files
    class_keys = pickle.load(
        open(os.path.join(path_indir, eventfiles[0]), 'rb'))[0].keys()
    print('class keys:', class_keys)

    # define dictionary
    event2word = {}
    word2event = {}

    corpus_kv = collections.defaultdict(list)
    for file in eventfiles:
        for event in pickle.load(open(
            os.path.join(path_indir, file), 'rb')):
            for key in class_keys:
                corpus_kv[key].append(event[key])

    for ckey in class_keys:
        class_unique_vals = sorted(
            set(corpus_kv[ckey]), key=lambda x: (not isinstance(x, int), x))
        event2word[ckey] = {key: i for i, key in enumerate(class_unique_vals)}
        word2event[ckey] = {i: key for i, key in enumerate(class_unique_vals)}

    # print
    print('[class size]')
    for key in class_keys:
        print(' > {:10s}: {}'.format(key, len(event2word[key])))

    # save
    path_dict = os.path.join(path_root, 'dictionary.pkl')
    pickle.dump((event2word, word2event), open(path_dict, 'wb'))

    # --- compile each --- #
    # reload
    event2word, word2event = pickle.load(open(path_dict, 'rb'))
    for fidx in range(len(eventfiles)):
        file = eventfiles[fidx]
        events_list = pickle.load(open(
            os.path.join(path_indir, file), 'rb'))
        fn = os.path.basename(file)
        path_outfile = os.path.join(path_outdir, fn)

        print('({}/{})'.format(fidx, len(eventfiles)))
        print(' > from:', file)
        print(' >   to:', path_outfile)

        words = []
        for eidx, e in enumerate(events_list):
            words_tmp = [
                event2word[k][e[k]] for k in class_keys
            ]
            words.append(words_tmp)

        # save
        path_outfile = os.path.join(path_outdir, file + '.npy')
        fn = os.path.basename(path_outfile)
        os.makedirs(path_outfile[:-len(fn)], exist_ok=True)
        np.save(path_outfile, words)


import json

from tracardi.domain.resources.elastic_resource_config import ElasticResourceConfig, ElasticCredentials
from tracardi.service.plugin.domain.register import Plugin, Spec, MetaData, Documentation, PortDoc, Form, FormGroup, \
    FormField, FormComponent
from tracardi.service.plugin.runner import ActionRunner
from tracardi.service.plugin.domain.result import Result
from tracardi.service.storage.elastic_client import ElasticClient
from .model.config import Config
from elasticsearch import ElasticsearchException
from tracardi.service.storage.driver.elastic import resource as resource_db
from tracardi.service.plugin.plugin_endpoint import PluginEndpoint

def validate(config: dict):
    config = Config(**config)

    # Try parsing JSON just for validation purposes
    try:
        json.loads(config.query)
    except json.JSONDecodeError as e:
        raise ValueError(str(e))

    return config

class Endpoint(PluginEndpoint):

    @staticmethod
    async def fetch_indices(config: dict):
        config = ElasticResourceConfig(**config)
        return await config.get_indices()

class ElasticSearchFetcher(ActionRunner):

    _client: ElasticClient
    config: Config

    async def set_up(self, init):
        config = Config(**init)
        resource = await resource_db.load(config.source.id)

        self.config = config
        credentials = resource.credentials.get_credentials(self, ElasticCredentials)
        self._client = credentials.get_client()

    async def run(self, payload: dict, in_edge=None) -> Result:

        try:
            query = json.loads(self.config.query)

            if 'size' not in query:
                query["size"] = 20

            if query["size"] > 50:
                self.console.warning("Fetching more then 50 records may impact the GUI performance.")

            result = await self._client.search(
                index=self.config.index.id,
                query=query
            )

            await self._client.close()

        except ElasticsearchException as e:
            self.console.error(str(e))
            return Result(port="error", value={
                "message": str(e)
            })

        return Result(port="result", value=result)

def register() -> Plugin:
    return Plugin(
        start=False,
        spec=Spec(
            module=__name__,
            className=ElasticSearchFetcher.__name__,
            inputs=["payload"],
            outputs=["result", "error"],
            version='0.6.0.1',
            license="MIT",
            author="Dawid Kruk",
            init={
                "source": {
                    "name": None,
                    "id": None
                },
                "index": None,
                "query": "{\"query\":{\"match_all\":{}}}"
            },
            manual="elasticsearch_query_action",
            form=Form(
                groups=[
                    FormGroup(
                        name="Plugin configuration",
                        fields=[
                            FormField(
                                id="source",
                                name="Elasticsearch resource",
                                description="Please select your Elasticsearch resource.",
                                component=FormComponent(type="resource", props={"label": "Resource",
                                                                                "tag": "elasticsearch"})
                            ),
                            FormField(
                                id="index",
                                name="Elasticsearch index",
                                description="Please select Elasticsearch index you want to search.",
                                component=FormComponent(type="autocomplete", props={
                                    "label": "Index",
                                    "endpoint": {
                                        "url": Endpoint.url(__name__, "fetch_indices"),
                                        "method": "post"
                                    }
                                })
                            ),
                            FormField(
                                id="query",
                                name="Query",
                                description="Please provide Elasticsearch DSL query.",
                                component=FormComponent(type="json", props={"label": "DSL query"})
                            )
                        ]
                    )
                ]
            )
        ),
        metadata=MetaData(
            name='Query Elasticsearch',
            desc='Gets data from given Elasticsearch resource',
            icon='elasticsearch',
            group=["Databases"],
            tags=['database', 'nosql'],
            documentation=Documentation(
                inputs={
                    "payload": PortDoc(desc="This port takes payload object.")
                },
                outputs={
                    "result": PortDoc(desc="This port returns result of querying ElasticSearch instance."),
                    "error": PortDoc(desc="This port returns error if one occurs, or if received result contains more "
                                          "than 20 records.")
                }
            )
        )
    )


from __future__ import print_function
import Pyro4
import person

Pyro4.config.SERIALIZER = 'pickle'
Pyro4.config.SERIALIZERS_ACCEPTED.add('pickle')
Pyro4.config.REQUIRE_EXPOSE = True
Pyro4.config.DETAILED_TRACEBACK = True

class Warehouse(object):
    def __init__(self):
        self.contents = ["chair", "bike", "flashlight", "laptop", "couch"]

    def list_contents(self):
        return self.contents

    def take(self, name, item):
        self.contents.remove(item)
        print("{0} took the {1}.".format(name, item))

    def store(self, name, item):
        self.contents.append(item)
        print("{0} stored the {1}.".format(name, item))

def main():
    warehouse = Warehouse()
    #Pyro4.config.HOST = "129.174.126.30:5555"
    daemon = Pyro4.Daemon()
    Pyro4.Daemon.serveSimple(
            {
                warehouse: "example.warehouse"
            },
	    daemon=daemon,
            ns = False)

if __name__=="__main__":
    main()


import numpy as np
from sklearn import svm
from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

def data_split(x, y, db):
    train_indexes = np.where(db.tr == 1)[0]
    test_indexs = [i for i in range(len(x)) if i not in train_indexes]
    x_train, y_train = x[train_indexes], y[train_indexes]
    x_test, y_test = x[test_indexs], y[test_indexs]

    return x_train, y_train, x_test, y_test

def plot_confusion_matrix(matrix, class_names):

    plt.figure(figsize=(6, 4))
    sns.heatmap(matrix, annot=True, fmt="d", cbar=False,
                xticklabels=class_names,
                yticklabels=class_names)
    plt.xlabel("Predicted labels")
    plt.ylabel("Real labels")
    plt.show()

def test_approach(x, db):
    y = db.label
    x_train, y_train, x_test, y_test = data_split(x, y, db)

    model = svm.SVC(C=20)
    model.fit(x_train, y_train)
    y_pred = model.predict(x_test)

    accuracy = accuracy_score(y_test, y_pred)
    print("Accuracy is: ", accuracy)

    # Calcola la matrice di confusione
    cf_matrix = confusion_matrix(y_test, y_pred)
    class_names = [db.id_class2name[id_label] for id_label in np.unique(y)]

    plot_confusion_matrix(cf_matrix, class_names)


from flask import Flask, render_template, url_for, request, redirect
import json
from datetime import datetime

app = Flask(__name__)
print(__name__)

@app.route("/")
def my_home():
    return render_template("index.html")

@app.route("/success.html")
def my_success():
    return render_template("success.html")

def save_to_db(data):
    print(f"Saving to db: {data}")
    return True

@app.route("/submit_form", methods=["POST"])
def submit_form():
    data_dict = request.form.to_dict()
    print(f"Received: {data_dict}")

    name = data_dict["name"]

    dob = data_dict["dob"]
    current_year = int(datetime.now().strftime("%Y"))
    dob_year = int(dob.split("-")[0])
    age = current_year - dob_year

    # Address: City State Country
    address = data_dict["address"]
    address_parts = address.split(" ")
    city = address_parts[0]
    state = address_parts[1]
    country = address_parts[2]
    location = f"{city},{state},{country}"

    save_data = {
        "name": name,
        "age": age,
        "location": location
    }

    saved = save_to_db(save_data)
    if saved:
       return redirect("success.html")


import matplotlib.pyplot as plt

def small_graph():
    small_graph_dict = dict(kivi={'visoki', 'viroza'},
                            viroza={'zarazni'},
                            zarazni={'niski', 'nitrati', 'niti'},
                            niski={'kivi'},
                            visoki={'kivi'},
                            nitrati={'tikovina'},
                            virulentni={'niski', 'nitrati'},
                            tikovina=set(),
                            niti={'tikovina'})
    graph = small_graph_dict
    dfs(small_graph_dict, 'kivi')
    # e = next(iter(graph))
    # print(e.startswith('Ki'.casefold()))
    # dfs(visited_set, graph, 'kivi')

def writing_ending_letters():
    ending_letters = set()
    with open('words_filled.csv', 'r') as file:
        reader = csv.reader(file, delimiter=',')
        for row in reader:
            # 0 index, 1 rijec, 2, pocetna slova, 3 zavrsna slova
            ending_letters.add(row[3])

    print(len(ending_letters))
    with open('ending_letters.txt', 'w') as f:
        f.write(str(ending_letters))

    ending_letters_dict = dict()
    for last_two_letters in ending_letters:
        with open('words_filled.csv', 'r') as file:
            reader = csv.reader(file, delimiter=',')
            set_of_connecting_words = set()
            for row in reader:
                if row[2] == last_two_letters:
                    set_of_connecting_words.add(row[1])
            ending_letters_dict[last_two_letters] = set_of_connecting_words

    print(ending_letters_dict)
    with open('ending_letters_dict.txt', 'w') as f:
        f.write(str(ending_letters_dict))

class Word:
    def __init__(self, text, visited: bool, ending_letters: string):
        self.connect_to_set = set()
        self.text = text
        self.visited = visited
        self.ending_letters = ending_letters

    def add_to_set(self, word_to_add):
        self.connect_to_set.add(word_to_add)

    def print_set_of_connecting_words(self):
        for w in self.connect_to_set:
            print(w.text)

def drawing_graph(list_of_words: [Word]):
    directed_graph = nx.DiGraph()
    for node in list_of_words:
        for edge in node.connect_to_set:
            directed_graph.add_edge(node.text, edge.text)
    nx.draw(directed_graph, with_labels=True)
    plt.show()
    # G.add_edges_from([('kivi', 'viroza'), ('kivi', 'visoki')])
    # print(G.nodes())
    # print(G.edges().data())

def reading_file_words_objects():
    with open('words_filled.csv', 'r') as file:
        reader = csv.reader(file, delimiter=',')
        words = []
        for row in itertools.islice(reader, 500):
            # print(row)
            # 0 index, 1 rijec, 2, pocetna slova, 3 zavrsna slova
            word = Word(row[1], False, row[3])
            words.append(word)
        for word_item in words:
            for w in words:
                if w.text.startswith(word_item.ending_letters):
                    word_item.add_to_set(w)

def playing_with_word_objects():
    kivi = Word('kivi', False, 'vi')
    visoki = Word('visoki', False, 'ki')
    viroza = Word('viroza', False, 'za')
    zarazni = Word('zarazni', False, 'ni')
    niski = Word('niski', False, 'ki')
    nitrati = Word('nitrati', False, 'ti')
    virulentni = Word('virulentni', False, 'ni')
    tikovina = Word('tikovina', False, 'na')

    graph_objects = {kivi: set([visoki, viroza]),
                     viroza: set([zarazni]),
                     zarazni: set([niski, nitrati]),
                     niski: set([kivi]),
                     visoki: set([kivi]),
                     nitrati: set([tikovina]),
                     virulentni: set([niski, nitrati]),
                     tikovina: set()}

def print_a_graph(visited: set(), graph, source):
    # This is not working properly
    print("==========")
    print(f"Source: {source}")
    visited_list.append(source)
    for node in graph:
        queue = [source]
        while len(queue) > 0:
            current = queue.pop()
            print(current)
            visited.add(current)
            for c in graph[current]:
                if c not in visited:
                    queue.append(c)


#!usr/bin/env python3
# coding=utf-8
# @IDE    :PyCharm
# @File   :vtk_save.py
# @Date   :2019/5/8

from vtk.util.vtkImageImportFromArray import *
import vtk
import SimpleITK as sitk
import cv2
import numpy as np

def getSeg(source, mask):
    source = np.reshape(source,(512,512))
    mask = np.reshape(mask,(512,512))*255
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    masks = cv2.erode(mask, kernel)
    _, contours, _ = cv2.findContours(masks, cv2.RETR_TREE,
                                      cv2.CHAIN_APPROX_SIMPLE)
    for i in range(512):
        for j in range(512):
            res = cv2.pointPolygonTest(contours[0], (i, j), False)
            if res == -1:
                source[j, i] = 0.0
    return source

v_path = r'F:\sliver_07\volume-0.nii'
s_path = r'F:\sliver_07\segmentation-0.nii'  # segmentation volume
ds = sitk.ReadImage(s_path)
s_data = sitk.GetArrayFromImage(ds)
s_data[s_data == 2] = 1
s_data = s_data[40:]
s_data = s_data.astype(np.uint8)
print(s_data.shape)
# print('------------read end segmentation---------------')
# dv = sitk.ReadImage(v_path)
# v_data = sitk.GetArrayFromImage(dv)
# print('------------read end volume---------------')
# result = []
# for i in range(v_data.shape[0]):
#     if np.sum(s_data[i] == 1) < 5:
#         data = np.zeros(shape=[512, 512], dtype=np.float32)
#         result.append(data)
#         continue
#     print(i)
#     data = getSeg(v_data[i],s_data[i])
#     result.append(data)
# result = np.array(result)
# print('------------seg end---------------')
spacing = ds.GetSpacing()
print(spacing)

img_arr = vtkImageImportFromArray()
img_arr.SetArray(s_data)
img_arr.SetDataSpacing(spacing)
origin = (0, 0, 0)
img_arr.SetDataOrigin(origin)
img_arr.Update()
srange = img_arr.GetOutput().GetScalarRange()
print(srange)
print(img_arr.GetOutput().GetSpacing())
min = srange[0]
max = srange[1]
diff = max - min
slope = 4000 / diff
inter = -slope * min
shift = inter / slope

shifter = vtk.vtkImageShiftScale()  # 对偏移和比例参数来对图像数据进行操作 数据转换
shifter.SetShift(shift)
shifter.SetScale(slope)
shifter.SetOutputScalarTypeToUnsignedShort()
shifter.SetInputData(img_arr.GetOutput())
shifter.ReleaseDataFlagOff()
shifter.Update()

print(shifter.GetOutput().GetSpacing())

skinExtractor = vtk.vtkContourFilter()
skinExtractor.SetInputConnection(shifter.GetOutputPort())
skinExtractor.SetValue(0,10)

# vtkWriter = vtk.vtkXMLImageDataWriter()
vtkWriter = vtk.vtkPolyDataWriter()
vtkWriter.SetInputConnection(skinExtractor.GetOutputPort())
vtkWriter.SetFileName('../result/seg.vtk')
vtkWriter.Write()
print('------------write end---------------')

# def StartInteraction():
#     renWin.SetDesiredUpdateRate(10)
#
# def EndInteraction():
#     renWin.SetDesiredUpdateRate(0.001)
#
# def ClipVolumeRender(obj):
#     obj.GetPlanes(planes)
#     volumeMapper.SetClippingPlanes(planes)
#
# ren = vtk.vtkRenderer()
# renWin= vtk.vtkRenderWindow()
# renWin.AddRenderer(ren)
# iren = vtk.vtkRenderWindowInteractor()
# iren.SetRenderWindow(renWin)
#
# tfun = vtk.vtkPiecewiseFunction()  # 梯度不透明度函数
# tfun.AddPoint(1000.0, 0)
# tfun.AddPoint(2300.0, 0.3)
# tfun.AddPoint(2501.0, 0.4)
# tfun.AddPoint(2600.0, 0.5)
# tfun.AddPoint(2700.0, 0.6)
# tfun.AddPoint(2900.0, 0.7)
# tfun.AddPoint(3024.0, 1)
#
#
# ctfun = vtk.vtkColorTransferFunction()  # 颜色传输函数
# ctfun.AddRGBPoint(0.0, 0.5, 0.0, 0.0)
# ctfun.AddRGBPoint(600.0, 1.0, 0.5, 0.5)
# ctfun.AddRGBPoint(1280.0, 0.9, 0.2, 0.3)
# ctfun.AddRGBPoint(1960.0, 0.81, 0.27, 0.1)
# ctfun.AddRGBPoint(2200.0, 0.9, 0.2, 0.3)
# ctfun.AddRGBPoint(2500.0, 1, 0.5, 0.5)
# ctfun.AddRGBPoint(3024.0, 0.5, 0.5, 0.5)
#
# volumeMapper = vtk.vtkGPUVolumeRayCastMapper()
# volumeMapper.SetInputData(shifter.GetOutput())
# volumeProperty = vtk.vtkVolumeProperty()
# volumeProperty.SetColor(ctfun)
# volumeProperty.SetScalarOpacity(tfun)
# volumeProperty.SetInterpolationTypeToLinear()
# volumeProperty.ShadeOn()
#
# newvol = vtk.vtkVolume()
# newvol.SetMapper(volumeMapper)
# newvol.SetProperty(volumeProperty)
#
# outline = vtk.vtkOutlineFilter()
# outline.SetInputConnection(shifter.GetOutputPort())
#
# outlineMapper = vtk.vtkPolyDataMapper()
# outlineMapper.SetInputConnection(outline.GetOutputPort())
#
# outlineActor = vtk.vtkActor()
# outlineActor.SetMapper(outlineMapper)
#
# ren.AddActor(outlineActor)
# ren.AddVolume(newvol)
# ren.SetBackground(0, 0, 0)
# renWin.SetSize(600, 600)
#
# planes = vtk.vtkPlanes()
#
# boxWidget = vtk.vtkBoxWidget()
# boxWidget.SetInteractor(iren)
# boxWidget.SetPlaceFactor(1.0)
# boxWidget.PlaceWidget(0,0,0,0,0,0)
# boxWidget.InsideOutOn()
# boxWidget.AddObserver("StartInteractionEvent", StartInteraction)
# boxWidget.AddObserver("InteractionEvent",  ClipVolumeRender)
# boxWidget.AddObserver("EndInteractionEvent",  EndInteraction)
#
# outlineProperty = boxWidget.GetOutlineProperty()
# outlineProperty.SetRepresentationToWireframe()
# outlineProperty.SetAmbient(1.0)
# outlineProperty.SetAmbientColor(1, 1, 1)
# outlineProperty.SetLineWidth(9)
#
# selectedOutlineProperty = boxWidget.GetSelectedOutlineProperty()
# selectedOutlineProperty.SetRepresentationToWireframe()
# selectedOutlineProperty.SetAmbient(1.0)
# selectedOutlineProperty.SetAmbientColor(1, 0, 0)
# selectedOutlineProperty.SetLineWidth(3)
#
# ren.ResetCamera()
# iren.Initialize()
# renWin.Render()
# iren.Start()


import setuptools

with open("./README.md", "r") as file:
    long_description = file.read()

setuptools.setup(
    name="httpcats",
    version="1.1.1",
    author="itsmewulf",
    author_email="wulf.developer@gmail.com",
    description="Get URLs to your favourite HTTP Cats easily!",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/itsmewulf/httpcats",
    keywords="cat, http, request, api, wrapper",
    packages=setuptools.find_packages(),
    include_package_data=True,
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    python_requires='>=3.6'
)


from django.contrib import admin

# import from apps here
from app_control.models import FunctionController

class FunctionControllerAdmin(admin.ModelAdmin):
    """
    功能开关表注册设置
    """

    list_display = ("func_code", "func_name", "enabled", "create_time", "func_developer")
    list_filter = ("func_code",)
    search_fields = ("func_code",)

admin.site.register(FunctionController, FunctionControllerAdmin)


print('='*12 + ' DESAFIO 41 ' + '='*12)
from datetime import date
ano = int(input('Digite seu ano de nascimento: '))
idade = date.today().year - ano
print('O atleta tem {} anos.'.format(idade))
if idade <= 9:
    classe = 'Mirim'
elif 9 < idade <= 14:
    classe = 'Infantil'
elif 14 < idade <= 19:
    classe = 'Junior'
elif 19 < idade <= 25:
    classe = 'Sênior'
else:
    classe = 'Master'
print('Classificação: \033[1m{}\033[m'.format(classe))


from setuptools import setup, find_packages

VERSION = '0.0.1'
DESCRIPTION = 'A package with useful functions relating to space.'

setup(
    name="Spacemath",
    version=VERSION,
    author="Periareion (Anton Sollman)",
    author_email="<periareion05@gmail.com>",
    description=DESCRIPTION,
    packages=find_packages(),
    install_requires=['numpy'],
    keywords=['space', 'math'],
    classifiers=[
        "Programming Language :: Python :: 3",
        "Operating System :: Unix",
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
    ]
)


import os
import traceback
from django.core.management.base import BaseCommand, CommandError, NoArgsCommand
from django.conf import settings
from azurepython3.blobservice import BlobService

class Command(NoArgsCommand):
    help = """Migrate from an existing FileSystemStorage to AzureStorage.
When executing the command "azuremigrate" it will scan for media files in MEDIA_ROOT
and upload all the files to the configured default container."""

    def handle_noargs(self, **options):
        # ensure that project has required configuration
        if not os.path.exists(settings.MEDIA_ROOT):
            raise CommandError('Cannot migrate files from non existing MEDIA_ROOT (%s)' % settings.MEDIA_ROOT)
        if not hasattr(settings, 'AZURE_ACCOUNT_NAME'):
            raise CommandError('AZURE_ACCOUNT_NAME setting missing')
        if not hasattr(settings, 'AZURE_ACCOUNT_KEY'):
            raise CommandError('AZURE_ACCOUNT_KEY setting missing')
        if not hasattr(settings, 'AZURE_DEFAULT_CONTAINER'):
            raise CommandError('AZURE_DEFAULT_CONTAINER settings missing')

        # get service interface
        service = BlobService(settings.AZURE_ACCOUNT_NAME, settings.AZURE_ACCOUNT_KEY)

        self.stdout.write('Starting migration from "%s" to '
                          'Cloud Storage container "%s"' % (settings.MEDIA_ROOT, settings.AZURE_DEFAULT_CONTAINER))

        for root, dirs, files in os.walk(settings.MEDIA_ROOT):
            for file in files:
                path = os.path.join(root, file)
                blobname = os.path.relpath(path, settings.MEDIA_ROOT).replace('\\', '/')
                self.stdout.write(blobname + "...", ending='')
                try:
                    with open(path, 'rb') as f:
                        if service.create_blob(settings.AZURE_DEFAULT_CONTAINER, blobname, bytearray(f.read())):
                            self.stdout.write('ok')
                        else:
                            self.stdout.write('fail')
                except Exception as e:
                    self.stdout.write('fail')
                    traceback.print_exc()
                    self.stdout.write('aborted migration.')
                    return

        self.stdout.write('migration complete')


import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from matplotlib import gridspec
import numpy as np
import random

from matplotlib import rc
# rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
# rc('font',**{'family':'serif','serif':['Palatino']})
rc('font',**{'family':'serif'})
rc('text', usetex=True)
rc('lines', linewidth=2)
rc('pdf',compression=0)
rc('figure', figsize= [6.4, 2.9], dpi=300, ) # figure size is in inches
rc('font',size=16)
rc('axes', grid=True)

saveFig = False
learnCount = 0

random.seed(1)
np.random.seed(1)
torch.manual_seed(1)

class LMS(torch.nn.Module):
    def __init__(self):
        super(LMS, self).__init__()
        self.neuron = torch.nn.Linear(10, 1)

    def forward(self, x):
        return self.neuron(x)

class nonlinear_LMS(torch.nn.Module):
    def __init__(self):
        super(nonlinear_LMS, self).__init__()
        self.neuron = torch.nn.Linear(10, 1)

    def forward(self, x):
        return 2 * torch.tanh(self.neuron(x))

class deep_network(torch.nn.Module):
    def __init__(self, layers):
        super(deep_network, self).__init__()
        self.layers = []
        for i in range(len(layers)):
            if i == 0:
                self.layers.append(torch.nn.Linear(10, layers[i]))
            else:
                self.layers.append(torch.nn.Linear(layers[i-1], layers[i]))

        self.layers = nn.ModuleList(self.layers)
    def forward(self, x):
        for i in range(len(self.layers)-1):
            x = F.relu(self.layers[i](x))
        return self.layers[-1](x)

def train_network(net, X, y, N, epochs, learning_rate):
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)

    train_loss = [0 for i in range(epochs)]
    test_loss = [0 for i in range(epochs)]

    for t in range(epochs):
        # Forward pass
        y_pred = net(X[:N, :])

        # Compute loss
        train_loss[t] = loss_fn(y_pred, y[:N, :])

        optimizer.zero_grad()

        # Backward pass.
        train_loss[t].backward()

        # Update
        optimizer.step()

        test_loss[t] = loss_fn(net(X[N:, :]), y[N:, :])

    return train_loss, test_loss

def plot_data(X, y):
    global learnCount
    plt.figure(figsize=(12, 8))
    gs = gridspec.GridSpec(2, 1)
    ax = plt.subplot(gs[0])

    for i in range(10):
        plt.plot(X[:, i], label='$x_{' + str(i + 1) + '}[n]$')

    plt.axvline(x=50, color='k')
    plt.text(0.25, 0.9, 'Train', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    plt.text(0.75, 0.9, 'Test', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    leg = plt.legend(bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
                     ncol=10, mode="expand", borderaxespad=0., fontsize=11)
    leg.get_frame().set_alpha(1)

    x_vect = ''.join(['x_{' + str(i + 1) + '}[n], ' for i in range(9)])

    plt.title('$\mathbf{x}[n] = [$ $' + x_vect + 'x_{10}[n]$ $]^T$', y=1.15, fontsize=14)
    plt.minorticks_on()
    plt.grid(b=True, which='major', color='k', linestyle='-', alpha=0.2)
    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.1)
    plt.tight_layout()

    ax = plt.subplot(gs[1])
    plt.plot(y)
    plt.axvline(x=50, color='k')
    plt.text(0.25, 0.9, 'Train', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    plt.text(0.75, 0.9, 'Test', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    leg = plt.legend(['$y[n]$'], fontsize=11)
    leg.get_frame().set_alpha(1)
    plt.title('$y[n] = \phi( \mathbf{x}[n] )$ where $\phi$ is highly non-linear', y=1.02, fontsize=14)

    plt.minorticks_on()
    plt.grid(b=True, which='major', color='k', linestyle='-', alpha=0.2)
    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.1)
    plt.tight_layout()

    if saveFig:
        plt.savefig('../figures/Q4_LMS2DL_4_7+8_%02i.pdf' % learnCount,format='pdf', bbox_inches="tight")
        learnCount = learnCount + 1

    plt.show()

def plot_output(X, y, models, deep_network_layers):
    plt.figure(figsize=(12,4))
    ax = plt.subplot(1,1,1)
    plt.plot(y)
    [plt.plot(model(torch.Tensor(X)).detach().numpy()) for model in models]
    plt.axvline(x=50,color='k')
    plt.text(0.25, 0.9, 'Train', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    plt.text(0.75, 0.9, 'Test', horizontalalignment='center', fontsize=14, transform = ax.transAxes)
    leg = plt.legend(['$y[n]$','Single Neuron (linear)','Single Neuron (tanh)','Deep Network ('+str(deep_network_layers)+', relu)'],
               bbox_to_anchor=(0., 1.02, 1., .102), loc=3,
               ncol=10, mode="expand", borderaxespad=0., fontsize=11)
    leg.get_frame().set_alpha(1)

    plt.minorticks_on()
    plt.grid(b=True, which='major', color='k', linestyle='-', alpha=0.2)
    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.1)
    plt.tight_layout()

    if saveFig:
        global learnCount
        plt.savefig('../figures/Q4_LMS2DL_4_7+8_%02i.pdf' % learnCount,format='pdf', bbox_inches="tight")
        learnCount = learnCount + 1

    plt.show()

def train_models(X, y, models, epochs=100, learning_rate=1e-2):
    loss = []

    for i, model in enumerate(models):
        (train_loss, test_loss) = train_network(model,
                                        X=torch.Tensor(X),
                                        y=torch.Tensor(y).reshape(-1,1),
                                        N=50,
                                        epochs=epochs,
                                        learning_rate=learning_rate)
        loss.append((train_loss, test_loss))

    return loss

def plot_learning_curves(loss):
    plt.figure(figsize=(12,4))
    gs = gridspec.GridSpec(1, 3)
    subPlotTitles = ['Single Neuron (linear)','Single Neuron (tanh)','Deep Network Loss']
    for i, l in enumerate(loss):
        ax = plt.subplot(gs[i])
        # plt.subplot(gs[i])
        ax.plot(l[0])
        ax.plot(l[1])
        ax.legend(['Train Loss', 'Test Loss'])
        ax.set_title(subPlotTitles[i])
        ax.minorticks_on()
        ax.grid(b=True, which='major', color='k', linestyle='-', alpha=0.2)
        ax.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.1)
        plt.tight_layout()
        plt.xlabel('Epoch')
        plt.ylabel('Error')
        # plt.subplots_adjust(top=0.8)  # Make room for the suptitle. (default: 0.9)

    # plt.suptitle('Learning Curves \n   ')

    if saveFig:
        global learnCount
        plt.savefig('../figures/Q4_LMS2DL_4_7+8_%02i.pdf' % learnCount,format='pdf', bbox_inches="tight")
        learnCount = learnCount + 1

    plt.show()

def savePlots(toggle):
    global saveFig
    global learnCount

    if toggle == True:
        learnCount = 0
        saveFig = True
    else:
        saveFig = False


#!/usr/bin/python3

import os
import sys
import re
import operator
from pathlib import Path

RECURSION_LIMIT = 9500
sys.setrecursionlimit(RECURSION_LIMIT + 10)

def count(word_list, ignore_words, word_freqs):
    if word_list == []:
        return
    else:
        word = word_list[0]
        if word not in ignore_words:
            if word in word_freqs:
                word_freqs[word] += 1
            else:
                word_freqs[word] = 1
        count(word_list[1:], ignore_words, word_freqs)

def wf_print(word_freqs):
    if word_freqs == []:
        return
    else:
        (w, c) = word_freqs[0]
        print(f'{w} - {c}')
        wf_print(word_freqs[1:])

ignore_words = set(open(os.path.join(Path.cwd(), 'res/ignore-words.txt')).read().split(','))
words = re.findall('[a-z]{2,}', open(sys.argv[1]).read().lower())
word_freqs = {}

for i in range(0, len(words), RECURSION_LIMIT):
    count(words[i:i + RECURSION_LIMIT], ignore_words, word_freqs)

wf_print(sorted(word_freqs.items(), key=operator.itemgetter(1), reverse=True)[:25])


from __future__ import print_function

import numpy as np
from PIL import Image
import os
import sys
import glob
import shutil

from argparser import args

train_examples_per_shape = args.n
test_examples_per_shape = max(1, train_examples_per_shape / 10)

image_w = 400
image_h = 400
shape_w = 50
shape_h = 50

max_x_noise = args.noise

max_y_noise = args.noise
#50 - 48%, 39%
#100 - 67$, 58%
#150 - 80%, 77%
#200 - 80%, 84%, 81%

hor_locations = {
    "left": args.e,#int(image_w / 4.0 - (shape_w / 2.0)),
    "right": image_w - shape_w - args.e,#int(3 * (image_w / 4.0) - (shape_w / 2.0)),
    "center": int(image_w / 2.0 - (shape_w / 2.0))
}
top_location = int(image_h / 2.0 - (shape_h / 2.0))

sim_id_t = 0

shapes = {}

def clean_up(output_dir):
    for f in glob.glob('{}/*'.format(output_dir)):
        shutil.rmtree(f)

def gen_stimulus(shape_no, location, output_dir, category, scramble = False):
    global sim_id_t, shapes, shape_h, shape_w

    img = Image.new(mode='RGB', size=(image_w, image_h), color=(255, 255, 255))
    if shape_no in shapes:
        shape = shapes[shape_no]
    else:
        shape = Image.open("{}/{}.png".format(shapes_dir, shape_no))
        #shape = shape.convert("L")
        #shape = ImageOps.invert(shape)
        shape.thumbnail((shape_h, shape_w))
        shapes[shape_no] = shape

    noise_x = int((max_x_noise / 2.0) - np.random.rand() * max_x_noise)
    noise_y = int((max_y_noise / 2.0) - np.random.rand() * max_y_noise)

    x = hor_locations[location] + noise_x
    y = top_location + noise_y
    img.paste(shape,box=(x, y))

    shape_output_dir = "{}/{}".format(output_dir, category)
    if not(os.path.exists(shape_output_dir)):
        os.makedirs(shape_output_dir)

    sim_id_t += 1

    stim_id = sim_id_t
    stim_file = "{}/{}.png".format(shape_output_dir, stim_id)
    img.save(stim_file)

#cleaning up trainng and test set
clean_up("./data/train/{}.{}".format(args.e, args.m))
clean_up("./data/test/{}.{}".format(args.e, args.m))

#Ryan's data, Leek
shapes_dir = "./stimuli/"
print("Generating train set (n = {}) ...".format(train_examples_per_shape), end='')
sys.stdout.flush()
for i in range(train_examples_per_shape):

    if (args.m == 1):
        for shape_no in range(1, 13):
            gen_stimulus(shape_no, "center", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
            gen_stimulus(shape_no + 12, "center", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)

    if (args.m == 2):
        for shape_no in range(1, 7):
            gen_stimulus(shape_no, "left", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
        for shape_no in range(7, 13):
            gen_stimulus(shape_no, "right", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
        for shape_no in range(13, 19):
            gen_stimulus(shape_no, "left", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)
        for shape_no in range(19, 25):
            gen_stimulus(shape_no, "right", "./data/train/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)

print("Done.")

print("Generating test set (n = {}) ...".format(test_examples_per_shape), end='')
sys.stdout.flush()
for _ in range(test_examples_per_shape):
    if (args.m == 1):
        for shape_no in range(1, 13):
            gen_stimulus(shape_no, "right", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
            gen_stimulus(shape_no + 12, "right", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)

    if (args.m == 2):
        for shape_no in range(1, 7):
            gen_stimulus(shape_no, "right", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
        for shape_no in range(7, 13):
            gen_stimulus(shape_no, "left", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 1)
        for shape_no in range(13, 19):
            gen_stimulus(shape_no, "right", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)
        for shape_no in range(19, 25):
            gen_stimulus(shape_no, "left", "./data/test/{}.{}.{}".format(args.e, args.m, args.noise), category = 2)

print("Done")


# A part of NonVisual Desktop Access (NVDA)
# This file is covered by the GNU General Public License.
# See the file COPYING for more details.
# Copyright (C) 2019 NV Access Limited

"""Unit tests for the orderedWinEventLimiter module.
"""
import inspect
import re
import unittest
from typing import List, Iterator, Callable
import winUser
from IAccessibleHandler import orderedWinEventLimiter
from IAccessibleHandler.orderedWinEventLimiter import OrderedWinEventLimiter

def softAssert(errorList: List[AssertionError], method: Callable, *args, **kwargs):
	try:
		method(*args, **kwargs)
	except AssertionError as e:
		errorList.append(e)

specialCaseEvents = [
	winUser.EVENT_SYSTEM_FOREGROUND,
	winUser.EVENT_OBJECT_FOCUS,
	winUser.EVENT_OBJECT_SHOW,
	winUser.EVENT_OBJECT_HIDE,
	winUser.EVENT_SYSTEM_MENUSTART,
	winUser.EVENT_SYSTEM_MENUEND,
	winUser.EVENT_SYSTEM_MENUPOPUPSTART,
	winUser.EVENT_SYSTEM_MENUPOPUPEND
]

def _getNonSpecialCaseEvents() -> Iterator[int]:
	objectOrSystemEvent = re.compile("^EVENT_(OBJECT|SYSTEM)_")
	for name, value in inspect.getmembers(winUser):
		if value not in specialCaseEvents and objectOrSystemEvent.match(name):
			yield value

nonSpecialCaseEvents: List[int] = list(_getNonSpecialCaseEvents())

class TestOrderedWinEventLimiter(unittest.TestCase):

	def test_nonSpecialCaseEvents(self):
		"""Test that the list of events without special cases matches expectations
		"""
		self.assertEqual(39, len(nonSpecialCaseEvents))

	def test_maxFocusEvents(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		for n in range(0, 5):
			limiter.addEvent(
				eventID=winUser.EVENT_OBJECT_FOCUS,
				window=n, objectID=n, childID=n, threadID=n,
			)
		events = limiter.flushEvents()
		windowIds = [
			e[1]  # window
			for e in events
		]
		expectedIds = [1, 2, 3, 4]
		self.assertEqual(expectedIds, windowIds)

	# Todo: Why do we get two events matching events?
	def test_foregroundOverwritesFocus(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		n = 1
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_FOCUS,
			window=n, objectID=n, childID=n, threadID=n,
		)
		limiter.addEvent(
			eventID=winUser.EVENT_SYSTEM_FOREGROUND,
			window=n, objectID=n, childID=n, threadID=n,
		)
		events = limiter.flushEvents()
		actualEvents = [(
			e[0],  # eventID
			e[1],  # window
		) for e in events
		]
		expectedEvents = [
			(winUser.EVENT_SYSTEM_FOREGROUND, n),
			(winUser.EVENT_SYSTEM_FOREGROUND, n)
		]
		self.assertEqual(expectedEvents, actualEvents)

	def test_showOverridesHide_whenDetailsMatch(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		n = 1
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_HIDE,
			window=n, objectID=n, childID=n, threadID=n,
		)
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_SHOW,
			window=n, objectID=n, childID=n, threadID=n,
		)
		events = limiter.flushEvents()
		actualEvents = [
			(e[0], e[1])
			for e in events
		]
		expectedEvents = [
			(winUser.EVENT_OBJECT_SHOW, n),
		]
		self.assertEqual(expectedEvents, actualEvents)

	def test_hideOverridesShow_whenDetailsMatch(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		n = 1
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_SHOW,
			window=n, objectID=n, childID=n, threadID=n,
		)
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_HIDE,
			window=n, objectID=n, childID=n, threadID=n,
		)
		events = limiter.flushEvents()
		actualEvents = [
			(e[0], e[1])
			for e in events
		]
		expectedEvents = [
			(winUser.EVENT_OBJECT_HIDE, n),
		]
		self.assertEqual(expectedEvents, actualEvents)

	def test_showNoOverridesHide_whenDetailsDontMatch(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		n = 1
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_HIDE,
			window=n, objectID=n, childID=n, threadID=n,
		)
		n = 2
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_SHOW,
			window=n, objectID=n, childID=n, threadID=n,
		)
		events = limiter.flushEvents()
		actualEvents = [
			(e[0], e[1])
			for e in events
		]
		expectedEvents = [
			(winUser.EVENT_OBJECT_HIDE, 1),
			(winUser.EVENT_OBJECT_SHOW, 2),
		]
		self.assertEqual(expectedEvents, actualEvents)

	def test_hideNoOverridesShow_whenDetailsDontMatch(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		n = 1
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_SHOW,
			window=n, objectID=n, childID=n, threadID=n,
		)
		n = 2
		limiter.addEvent(
			eventID=winUser.EVENT_OBJECT_HIDE,
			window=n, objectID=n, childID=n, threadID=n,
		)
		events = limiter.flushEvents()
		actualEvents = [
			(e[0], e[1])
			for e in events
		]
		expectedEvents = [
			(winUser.EVENT_OBJECT_SHOW, 1),
			(winUser.EVENT_OBJECT_HIDE, 2),
		]
		self.assertEqual(expectedEvents, actualEvents)

	def test_alwaysAllowedObjects_specialCaseEvents(self):
		# We have events from two unique objects:
		# Window, objectID, childID
		allowedSource = (1, 1, 1)

		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		for n in range(2000):  # send many events, to saturate all limits.
			eventId = specialCaseEvents[n % len(specialCaseEvents)]
			limiter.addEvent(eventId, *allowedSource, threadID=0)
		events = limiter.flushEvents(alwaysAllowedObjects=[allowedSource, ])

		expected = [
			# Two Foreground events, because they are added to multiple queues.
			# Added to both _focusEventCache and _genericEventCache
			# See also test_limitEventsPerThread
			(winUser.EVENT_SYSTEM_FOREGROUND, *allowedSource),
			(winUser.EVENT_SYSTEM_FOREGROUND, *allowedSource),
			(winUser.EVENT_OBJECT_FOCUS, *allowedSource),
			(winUser.EVENT_OBJECT_HIDE, *allowedSource),  # latest out of show / hide is kept
			(winUser.EVENT_SYSTEM_MENUPOPUPEND, *allowedSource),  # Only one menu event is allowed
		]
		self.assertEqual(expected, events)

	def test_alwaysAllowedObjects_onlyLatestEventKept(self):
		# We have events from two unique objects:
		# Window, objectID, childID
		allowedSource = (1, 1, 1)
		otherSource = (2, 2, 2,)

		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		for n in range(50):  # send many value changed events
			limiter.addEvent(winUser.EVENT_OBJECT_VALUECHANGE, *allowedSource, threadID=0)
			limiter.addEvent(winUser.EVENT_OBJECT_VALUECHANGE, *otherSource, threadID=0)
		events = limiter.flushEvents(alwaysAllowedObjects=[allowedSource, ])
		# only the most recent event of each object is kept, all previous duplicates are discarded
		self.assertEqual(2, len(events))

	def test_threadLimit_singleObject(self):
		"""Test that only the latest events are kept when the thread limit is exceeded
		"""
		# We have events from two unique objects:
		# Window, objectID, childID
		source = (2, 2, 2,)

		limiter = OrderedWinEventLimiter(maxFocusItems=4)

		for n in range(500):  # exceed the limit for a single thread
			eventId = nonSpecialCaseEvents[n % len(nonSpecialCaseEvents)]
			# same thread, different object. Use a second object to aid tracking.
			limiter.addEvent(eventId, *source, threadID=0)

		events = limiter.flushEvents()
		expectedEventCount = orderedWinEventLimiter.MAX_WINEVENTS_PER_THREAD
		self.assertEqual(expectedEventCount, len(events))

	def test_threadLimit_noCanary(self):
		"""Test that only the latest events are kept when the thread limit is exceeded
		"""
		limiter = OrderedWinEventLimiter(maxFocusItems=4)

		for n in range(500):  # exceed the limit for a single thread
			eventId = nonSpecialCaseEvents[n % len(nonSpecialCaseEvents)]
			# same thread, different object. Ensure there are no duplicates
			# Window, objectID, childID
			source = (2, 2, n,)
			limiter.addEvent(eventId, *source, threadID=0)

		events = limiter.flushEvents()

		errors = []
		expectedEventCount = orderedWinEventLimiter.MAX_WINEVENTS_PER_THREAD
		softAssert(errors, self.assertEqual, expectedEventCount, len(events))  # Fails with 11 actual events
		self.assertListEqual([], errors)

	def test_threadLimit_withCanaryAtStart(self):
		"""Test that only the latest events are kept when the thread limit is exceeded
		"""
		limiter = OrderedWinEventLimiter(maxFocusItems=4)

		# Window, objectID, childID
		canaryObject = (1, 1, 1)
		eventStartCanary = (winUser.EVENT_OBJECT_VALUECHANGE, *canaryObject)
		limiter.addEvent(*eventStartCanary, threadID=0)

		for n in range(500):  # exceed the limit for a single thread
			eventId = nonSpecialCaseEvents[n % len(nonSpecialCaseEvents)]
			# same thread, different object. Ensure there are no duplicates
			# Window, objectID, childID
			source = (2, 2, n,)
			limiter.addEvent(eventId, *source, threadID=0)

		events = limiter.flushEvents()

		errors = []
		expectedEventCount = orderedWinEventLimiter.MAX_WINEVENTS_PER_THREAD
		softAssert(errors, self.assertEqual, expectedEventCount, len(events))  # Fails with 11 actual events
		softAssert(errors, self.assertNotIn, eventStartCanary, events)
		self.assertListEqual([], errors)

	def test_threadLimit_canaryStartAndEnd(self):
		"""Test that only the latest events are kept when the thread limit is exceeded
		"""
		limiter = OrderedWinEventLimiter(maxFocusItems=4)

		# Window, objectID, childID
		canaryObject = (1, 1, 1)
		eventStartCanary = (winUser.EVENT_OBJECT_VALUECHANGE, *canaryObject)
		limiter.addEvent(*eventStartCanary, threadID=0)

		for n in range(500):  # exceed the limit for a single thread
			eventId = nonSpecialCaseEvents[n % len(nonSpecialCaseEvents)]
			# same thread, different object. Ensure there are no duplicates
			# Window, objectID, childID
			source = (2, 2, n,)
			limiter.addEvent(eventId, *source, threadID=0)

		# Note event type must differ from start canary to ensure they are not duplicates
		eventEndCanary = (winUser.EVENT_OBJECT_NAMECHANGE, *canaryObject)
		limiter.addEvent(*eventEndCanary, threadID=0)

		events = limiter.flushEvents()
		errors = []
		expectedEventCount = orderedWinEventLimiter.MAX_WINEVENTS_PER_THREAD
		softAssert(errors, self.assertEqual, expectedEventCount, len(events))  # Fails with 11 actual events
		softAssert(errors, self.assertIn, eventEndCanary, events)
		softAssert(errors, self.assertNotIn, eventStartCanary, events)
		self.assertListEqual([], errors)

	def test_alwaysAllowedObjects(self):
		"""Matches test_threadLimit_canaryStartAndEnd, but allows events from the first object
		"""
		limiter = OrderedWinEventLimiter(maxFocusItems=4)

		# Window, objectID, childID
		canaryObject = (1, 1, 1)
		eventStartCanary = (winUser.EVENT_OBJECT_VALUECHANGE, *canaryObject)
		limiter.addEvent(*eventStartCanary, threadID=0)

		for n in range(orderedWinEventLimiter.MAX_WINEVENTS_PER_THREAD):  # exceed the limit for a single thread
			eventId = nonSpecialCaseEvents[n % len(nonSpecialCaseEvents)]
			# same thread, different object. Ensure there are no duplicates
			# Window, objectID, childID
			source = (2, 2, n,)
			limiter.addEvent(eventId, *source, threadID=0)

		eventEndCanary = (winUser.EVENT_OBJECT_NAMECHANGE, *canaryObject)
		limiter.addEvent(*eventEndCanary, threadID=0)

		events = limiter.flushEvents(alwaysAllowedObjects=[canaryObject, ])
		# only the most recent event of each object is kept, all previous duplicates are discarded
		self.assertEqual(11, len(events))
		self.assertIn(eventStartCanary, events)
		self.assertEqual(eventStartCanary, events[0])
		self.assertIn(eventEndCanary, events)

	def test_limitEventsPerThread(self):
		limiter = OrderedWinEventLimiter(maxFocusItems=4)
		for n in reversed(range(2000)):  # send many events, to saturate all limits.
			limiter.addEvent(
				eventID=specialCaseEvents[n % len(specialCaseEvents)],
				window=n, objectID=n, childID=n,
				threadID=0,  # all events for same thread
			)
		events = limiter.flushEvents()
		windowIds = [
			e[1]  # window
			for e in events
		]
		# TODO: Note: repeated Id's (0 and 8) are EVENT_SYSTEM_FOREGROUND see test_maxFocusEvents
		expectedIds = [24, 19, 18, 16, 11, 10, 9, 8, 8, 4, 3, 2, 1, 0, 0]
		self.assertEqual(expectedIds, windowIds)
		#  equal to MAX_WINEVENTS_PER_THREAD=10
		# Plus 4 focus events,
		# Plus the last menu event.
		#  All totalling 15.
		self.assertEqual(len(windowIds), 15)


import string
from copy import copy
from random import sample

import numpy as np
import pytest

import filter_functions as ff
from filter_functions import numeric
from tests import testutil
from tests.testutil import rng

from . import qutip

plotting = pytest.importorskip('filter_functions.plotting',
                               reason='Skipping plotting tests for build without matplotlib')
if plotting is not None:
    import matplotlib.pyplot as plt
    from matplotlib import cycler

simple_pulse = testutil.rand_pulse_sequence(2, 1, 1, 1, btype='Pauli')
complicated_pulse = testutil.rand_pulse_sequence(2, 100, 3, 3)
two_qubit_pulse = testutil.rand_pulse_sequence(4, 1, 1, 6, btype='Pauli')

class PlottingTest(testutil.TestCase):

    def test_plot_pulse_train(self):
        # Call with default args
        fig, ax, leg = plotting.plot_pulse_train(simple_pulse)

        # Call with no axes but figure
        fig = plt.figure()
        fig, ax, leg = plotting.plot_pulse_train(simple_pulse, fig=fig)

        # Call with axes but no figure
        fig, ax, leg = plotting.plot_pulse_train(simple_pulse, axes=ax)

        # Call with custom args
        c_oper_identifiers = sample(
            complicated_pulse.c_oper_identifiers.tolist(), rng.integers(2, 4)
        )

        fig, ax = plt.subplots()
        fig, ax, leg = plotting.plot_pulse_train(complicated_pulse,
                                                 c_oper_identifiers,
                                                 fig=fig, axes=ax)

        # Test cycler arg
        cycle = cycler(color=['r', 'g', 'b'])
        fig, ax, leg = plotting.plot_pulse_train(simple_pulse, cycler=cycle)

        # invalid identifier
        with self.assertRaises(ValueError):
            plotting.plot_pulse_train(complicated_pulse,
                                      c_oper_identifiers=['foo'],
                                      fig=fig, axes=ax)

        # Test various keyword args for matplotlib
        plot_kw = {'linewidth': 1}
        subplot_kw = {'facecolor': 'r'}
        gridspec_kw = {'hspace': 0.2,
                       'wspace': 0.1}
        figure_kw = {'num': 1}

        fig, ax, leg = plotting.plot_pulse_train(simple_pulse, plot_kw=plot_kw,
                                                 subplot_kw=subplot_kw,
                                                 gridspec_kw=gridspec_kw,
                                                 **figure_kw)

        plt.close('all')

    def test_plot_filter_function(self):
        # Call with default args
        simple_pulse.cleanup('all')
        fig, ax, leg = plotting.plot_filter_function(simple_pulse)

        # Call with no axes but figure
        fig = plt.figure()
        fig, ax, leg = plotting.plot_filter_function(simple_pulse, fig=fig)

        # Call with axes but no figure
        fig, ax, leg = plotting.plot_filter_function(simple_pulse, axes=ax)

        # Non-default args
        n_oper_identifiers = sample(
            complicated_pulse.n_oper_identifiers.tolist(), rng.integers(2, 4)
        )

        fig, ax = plt.subplots()
        omega = ff.util.get_sample_frequencies(complicated_pulse, n_samples=50,
                                               spacing='log')
        fig, ax, leg = plotting.plot_filter_function(
            complicated_pulse, omega=omega,
            n_oper_identifiers=n_oper_identifiers,
            fig=fig, axes=ax, omega_in_units_of_tau=False
        )

        # Test cycler arg
        cycle = cycler(color=['r', 'g', 'b'])
        fig, ax, leg = plotting.plot_filter_function(simple_pulse, cycler=cycle)

        # invalid identifier
        with self.assertRaises(ValueError):
            plotting.plot_filter_function(complicated_pulse,
                                          n_oper_identifiers=['foo'],
                                          fig=fig, axes=ax)

        # Test different axis scales
        scales = ('linear', 'log')
        for xscale in scales:
            for yscale in scales:
                fig, ax, leg = plotting.plot_filter_function(simple_pulse,
                                                             xscale=xscale,
                                                             yscale=yscale)

        # Test various keyword args for matplotlib
        plot_kw = {'linewidth': 1}
        subplot_kw = {'facecolor': 'r'}
        gridspec_kw = {'hspace': 0.2,
                       'wspace': 0.1}
        figure_kw = {'num': 1}
        fig, ax, leg = plotting.plot_filter_function(simple_pulse,
                                                     plot_kw=plot_kw,
                                                     subplot_kw=subplot_kw,
                                                     gridspec_kw=gridspec_kw,
                                                     **figure_kw)

        plt.close('all')

    @pytest.mark.filterwarnings('ignore::UserWarning')  # tight_layout warning
    def test_plot_pulse_correlation_filter_function(self):
        omega = np.linspace(-1, 1, 50)
        concatenated_simple_pulse = ff.concatenate(
            (simple_pulse, simple_pulse),
            calc_pulse_correlation_FF=True,
            omega=omega
        )
        concatenated_complicated_pulse = ff.concatenate(
            (complicated_pulse, complicated_pulse),
            calc_pulse_correlation_FF=True,
            omega=omega
        )

        # Exception if not pulse correl. FF is cached
        with self.assertRaises(ff.util.CalculationError):
            plotting.plot_pulse_correlation_filter_function(simple_pulse)

        # Call with default args
        fig, ax, leg = plotting.plot_pulse_correlation_filter_function(
            concatenated_simple_pulse
        )

        # Non-default args
        n_oper_identifiers = sample(
            complicated_pulse.n_oper_identifiers.tolist(), rng.integers(2, 4)
        )

        fig, ax = plt.subplots()
        omega = np.linspace(-10, 10, 50)
        fig, ax, leg = plotting.plot_pulse_correlation_filter_function(
            concatenated_complicated_pulse, omega=omega,
            n_oper_identifiers=n_oper_identifiers, fig=fig,
            omega_in_units_of_tau=False
        )

        # Test cycler arg
        cycle = cycler(color=['r', 'g', 'b'])
        fig, ax, leg = plotting.plot_pulse_correlation_filter_function(concatenated_simple_pulse,
                                                                       cycler=cycle)

        # invalid identifiers
        with self.assertRaises(ValueError):
            plotting.plot_pulse_correlation_filter_function(
                concatenated_complicated_pulse,
                n_oper_identifiers=['foo'], fig=fig, axes=ax
            )

        # Test different axis scales
        scales = ('linear', 'log')
        for xscale in scales:
            for yscale in scales:
                fig, ax, leg = plotting.plot_pulse_correlation_filter_function(
                    concatenated_simple_pulse, xscale=xscale, yscale=yscale
                )

        # Test various keyword args for matplotlib
        plot_kw = {'linewidth': 1}
        subplot_kw = {'facecolor': 'r'}
        gridspec_kw = {'hspace': 0.2,
                       'wspace': 0.1}
        figure_kw = {'num': 1}
        fig, ax, leg = plotting.plot_pulse_correlation_filter_function(
            concatenated_simple_pulse, plot_kw=plot_kw, subplot_kw=subplot_kw,
            gridspec_kw=gridspec_kw, **figure_kw
        )

        plt.close('all')

    # Ignore deprecation warning caused by qutip
    @pytest.mark.filterwarnings('ignore::PendingDeprecationWarning')
    def test_plot_cumulant_function(self):
        omega = ff.util.get_sample_frequencies(simple_pulse)
        spectrum = 1e-4*np.sin(omega)/omega

        # Test calling with pulse, spectrum, omega
        fig, grid = plotting.plot_cumulant_function(simple_pulse, spectrum, omega,
                                                    colorscale='linear')
        fig, grid = plotting.plot_cumulant_function(simple_pulse, spectrum, omega, fig=fig)
        fig, grid = plotting.plot_cumulant_function(simple_pulse, spectrum, omega, grid=grid)

        # Test calling with precomputed transfer matrix
        K = numeric.calculate_cumulant_function(simple_pulse, spectrum, omega)
        fig, grid = plotting.plot_cumulant_function(cumulant_function=K)

        # Test calling with precomputed transfer matrix and pulse
        K = numeric.calculate_cumulant_function(simple_pulse, spectrum, omega)
        fig, grid = plotting.plot_cumulant_function(simple_pulse, cumulant_function=K)

        # Test calling with precomputed transfer matrix of ndim == 2
        K = numeric.calculate_cumulant_function(simple_pulse, spectrum, omega)
        fig, grid = plotting.plot_cumulant_function(cumulant_function=K[0])

        # Log colorscale
        fig, grid = plotting.plot_cumulant_function(cumulant_function=K, colorscale='log')

        # Non-default args
        n_oper_inds = sample(range(len(complicated_pulse.n_opers)), rng.integers(2, 4))
        n_oper_identifiers = complicated_pulse.n_oper_identifiers[n_oper_inds]

        basis_labels = []
        for i in range(4):
            basis_labels.append(string.ascii_uppercase[rng.integers(0, 26)])

        omega = ff.util.get_sample_frequencies(complicated_pulse, n_samples=50, spacing='log')
        spectrum = np.exp(-omega**2)
        K = numeric.calculate_cumulant_function(complicated_pulse, spectrum, omega)
        fig, grid = plotting.plot_cumulant_function(
            complicated_pulse, spectrum=spectrum, omega=omega,
            n_oper_identifiers=n_oper_identifiers, basis_labels=np.array(basis_labels),
            basis_labelsize=4, linthresh=1e-4, cmap=plt.cm.jet
        )
        fig, grid = plotting.plot_cumulant_function(
            cumulant_function=K[n_oper_inds], n_oper_identifiers=n_oper_identifiers,
            basis_labels=basis_labels, basis_labelsize=4, linthresh=1e-4,
            cmap=plt.cm.jet
        )

        # neither K nor all of pulse, spectrum, omega given
        with self.assertRaises(ValueError):
            plotting.plot_cumulant_function(complicated_pulse, spectrum)

        # invalid identifiers
        with self.assertRaises(ValueError):
            plotting.plot_cumulant_function(complicated_pulse, spectrum, omega,
                                            n_oper_identifiers=['foo'])

        # K and identifiers given but unequal length
        with self.assertRaises(ValueError):
            plotting.plot_cumulant_function(cumulant_function=K,
                                            n_oper_identifiers=n_oper_identifiers[0])

        # number of basis_labels not correct
        with self.assertRaises(ValueError):
            plotting.plot_cumulant_function(complicated_pulse, spectrum, omega,
                                            basis_labels=basis_labels[:2])

        # grid too small
        with self.assertRaises(ValueError):
            plotting.plot_cumulant_function(complicated_pulse, spectrum, omega,
                                            grid=grid[:1])

        # Test various keyword args for matplotlib for the two-qubit pulse
        spectrum = np.tile(spectrum, (6, 6, 1))
        grid_kw = {'axes_pad': 0.1}
        cbar_kw = {'orientation': 'horizontal'}
        imshow_kw = {'interpolation': 'bilinear'}
        figure_kw = {'num': 1}
        fig, ax = plotting.plot_cumulant_function(two_qubit_pulse, spectrum, omega,
                                                  imshow_kw=imshow_kw, grid_kw=grid_kw,
                                                  cbar_kw=cbar_kw, **figure_kw)

        plt.close('all')

    def test_plot_infidelity_convergence(self):
        n, infids = ff.infidelity(simple_pulse, lambda x: x**0, {}, test_convergence=True)
        fig, ax = plotting.plot_infidelity_convergence(n, infids)

        fig, ax = plt.subplots(1, 2)
        cfig, ax = plotting.plot_infidelity_convergence(n, infids, ax)
        self.assertIs(cfig, fig)

class LaTeXRenderingTest(testutil.TestCase):

    def test_plot_filter_function(self):
        pulse = copy(simple_pulse)
        pulse.c_oper_identifiers = np.array([f'B_{i}' for i in range(len(pulse.c_opers))])
        pulse.n_oper_identifiers = np.array([f'B_{i}' for i in range(len(pulse.n_opers))])
        with plt.rc_context(rc={'text.usetex': True}):
            _ = plotting.plot_pulse_train(pulse)
            _ = plotting.plot_filter_function(pulse)

@pytest.mark.skipif(
    qutip is None,
    reason='Skipping bloch sphere visualization tests for build without qutip')
class BlochSphereVisualizationTest(testutil.TestCase):

    def test_get_bloch_vector(self):
        states = [qutip.rand_ket(2) for _ in range(10)]
        bloch_vectors_qutip = plotting.get_bloch_vector(states)
        bloch_vectors_np = plotting.get_bloch_vector([state.full()
                                                      for state in states])

        for bv_qutip, bv_np in zip(bloch_vectors_qutip, bloch_vectors_np):
            self.assertArrayAlmostEqual(bv_qutip, bv_np)

    def test_get_states_from_prop(self):
        P = testutil.rand_unit(2, 10)
        Q = np.empty((11, 2, 2), dtype=complex)
        Q[0] = np.identity(2)
        for i in range(10):
            Q[i+1] = P[i] @ Q[i]

        psi0 = qutip.rand_ket(2)
        psi0_np = psi0.full()
        states = plotting.get_states_from_prop(Q, psi0)
        states_np = plotting.get_states_from_prop(Q, psi0_np)
        states_0 = plotting.get_states_from_prop(Q)
        self.assertArrayAlmostEqual(Q @ psi0_np, states)
        self.assertArrayAlmostEqual(Q @ psi0_np, states_np)
        self.assertArrayAlmostEqual(Q @ qutip.basis(2, 0).full(), states_0)

        with self.assertRaises(ValueError):
            plotting.get_states_from_prop(Q, rng.standard_normal((3, 1, 2)))
        with self.assertRaises(ValueError):
            plotting.get_states_from_prop(Q, rng.standard_normal((3, 4)))

    def test_plot_bloch_vector_evolution(self):
        # Call with default args
        b = plotting.plot_bloch_vector_evolution(simple_pulse)
        # Call with custom args
        b = plotting.init_bloch_sphere(background=True)
        b = plotting.plot_bloch_vector_evolution(simple_pulse,
                                                 psi0=qutip.basis(2, 1), b=b,
                                                 n_samples=50, show=False,
                                                 return_Bloch=True)

        b = plotting.plot_bloch_vector_evolution(complicated_pulse)

        # Test add_cbar kwarg
        b = plotting.plot_bloch_vector_evolution(simple_pulse, cmap='viridis', add_cbar=True)

        # Check exceptions being raised
        with self.assertRaises(ValueError):
            plotting.plot_bloch_vector_evolution(two_qubit_pulse)

        plt.close('all')

    def test_box_aspect(self):
        """Fix https://github.com/qutech/filter_functions/issues/41"""
        b = plotting.plot_bloch_vector_evolution(simple_pulse, return_Bloch=True)
        aspect = b.axes.get_box_aspect()
        self.assertAlmostEqual(aspect.std(), 0)