# -*- coding: utf-8 -*-
"""
Created on Mon Aug 17 19:21:08 2020

@author: bmcma
"""

import rnn
import numpy as np
import torch
import matplotlib.pyplot as plt
import svcca.cca_core as cca_core
from sklearn.manifold import MDS
from sklearn.cluster import KMeans
from sklearn.mixture import GaussianMixture
from sklearn.metrics import silhouette_score
import pdb
import argparse
from FP_Analysis import ComputeDistance

# TODO: refactor code as class
# class NetworkDistances:
#     # public:
#     def getDistances():
#         self._getActivations()
#     # private:
#     def loadRNNS():  # loads RNNs to be analyzed
#         pass
#     def _getActivations():  # gets the activations of all RNNs
#         pass

def getNumSVs(singularValues):
    sv_squares = singularValues**2
    svs_needed = 0
    explained_var = 0
    sv_square_totals = np.sum(sv_squares)
    while (explained_var < 0.95):
        explained_var += (singularValues[svs_needed]**2 / sv_square_totals)
        svs_needed+=1
    return svs_needed

def get_rdm_inputs(inputSize):
    # returns inputs for the RDM task
    #assert False, "you are using RDM, use N=1 network instead"
    NUM_INPUT_CONDITIONS = 10
    static_inputs = np.linspace(-0.1857, 0.1857, NUM_INPUT_CONDITIONS).reshape(1, NUM_INPUT_CONDITIONS)
    static_inputs = np.matmul(np.ones((750, 1)), static_inputs)
    static_inputs = torch.tensor(static_inputs).float()
    static_inputs = torch.unsqueeze(static_inputs.t(), 1)
    return static_inputs

def get_static_inputs(inputSize):
    # will get inputs for the N = 1 context task
    # N_INPUTS = 120   # needs to be a multiple of all possible N
    N_INPUTS_PER_CONTEXT = 10
    N_CONTEXTS = int(inputSize / 2)
    #N_INPUTS_PER_CONTEXT = int(N_INPUTS / N_CONTEXTS)
    MIN_INPUT = -0.1857
    MAX_INPUT = 0.1857
    static_inpt_arr = np.zeros((N_INPUTS_PER_CONTEXT*N_CONTEXTS, inputSize))
    #static_inpt_arr = np.zeros((N_INPUTS_PER_CONTEXT*N_CONTEXTS, inputSize))
    # loop over contexts
    for context in range(N_CONTEXTS):
       strt_ix = context*N_INPUTS_PER_CONTEXT
       end_ix = (context+1)*N_INPUTS_PER_CONTEXT
       static_inpt_arr[strt_ix:end_ix, context] = np.linspace(MIN_INPUT, MAX_INPUT, N_INPUTS_PER_CONTEXT)  # set inputs
       static_inpt_arr[strt_ix:end_ix, context+N_CONTEXTS] = 1  # set GO signal
    # static_inpt_arr now filled in with shape (20, 8)
    static_inpt_arr = np.expand_dims(static_inpt_arr, -1)  # (20, 8, 1)
    assert static_inpt_arr.shape == (N_INPUTS_PER_CONTEXT*N_CONTEXTS, inputSize, 1)
    inpts = np.matmul(static_inpt_arr, np.ones((1, 750))) # (20, 8, 750)
    # feed batch of 20 inpts into network to get ativations
    # cast static inputs to PyTorch Tensor
    static_inputs = torch.tensor(inpts).float() # 20, 8, 750
    return static_inputs

def getActivations(rnn_model):
    '''generates inputs for SVCCA'''
    switcher = {
            1:get_rdm_inputs,
            2:get_static_inputs,
            4:get_static_inputs,
            6:get_static_inputs,
            8:get_static_inputs,
            10:get_static_inputs,
            12:get_static_inputs
        }
    static_inputs = switcher[rnn_model._inputSize](rnn_model._inputSize)
    _, activations = rnn_model.feed(static_inputs, return_hidden=True)
    if rnn_model._useHeb:
        heb_mask = np.ones(50).astype(int)
        heb_mask[1] = 0; heb_mask[10] = 0; heb_mask[11] = 0
        activations = activations[:, np.where(heb_mask==1)[0], :]
        activations = np.swapaxes(activations, 0, 1).reshape(rnn_model._hiddenSize-3, -1)
        return activations #finalActivations    # 50 x 500  ----> numHidden x numInputs
    #finalActivations = activations[-1,:,:]       # keeps activation at end of trial only
    activations = np.swapaxes(activations, 0, 1).reshape(rnn_model._hiddenSize, -1)
    return activations #finalActivations    # 50 x 500  ----> numHidden x numInputs

# get a file of models to analyze
parser = argparse.ArgumentParser(description="Clusters RNNs by SVCCA")
parser.add_argument("fname", help="name of file containing RNNs to analyze")
parser.add_argument("-reload", action="store_true", default=False)
args = parser.parse_args()
file_of_rnns = open("models/"+args.fname, 'r')
models = [(line.strip()).split() for line in file_of_rnns]
file_of_rnns.close()

modelActivations = []
NUM_DIMENSIONS = 2

#models = [bptt, ga]
numModelsOfType = {}
count = 0
for i in range(len(models)):
    numModelsOfType[models[i][0]] = len(models[i])-1
    count += len(models[i]) - 1
numModelsOfType["total"] = count

if args.reload:
    distances = np.zeros((numModelsOfType["total"], numModelsOfType["total"]))
    for modelType in models:
        for count, model in enumerate(modelType):
            if (count == 0):
                print("model name:", model)
            else:
                print("name:", model)
                rnn_inst = rnn.loadRNN("models/"+model)
                activations = getActivations(rnn_inst)
                modelActivations.append(activations)

    # compute the distance matrix
    print("computing distance matrix ...\n")
    for i in range(len(modelActivations)):
        print("    Computing row", int(i))
        for j in range(i, len(modelActivations)):   # ~2X speedup using symmetry
            activationsI = modelActivations[i]
            activationsJ = modelActivations[j]
            activationsI -= np.mean(activationsI, axis=1, keepdims=True)
            activationsJ -= np.mean(activationsJ, axis=1, keepdims=True)

            # take top 20 singular values
            U1, s1, V1 = np.linalg.svd(activationsI, full_matrices=False)
            U2, s2, V2 = np.linalg.svd(activationsJ, full_matrices=False)

            #print("Number of singular values needed to explain 95% of variance:", getNumSVs(s1))
            NUM_DIMENSIONS = getNumSVs(s1)
            svacts1 = np.dot(s1[:NUM_DIMENSIONS]*np.eye(NUM_DIMENSIONS), V1[:NUM_DIMENSIONS])
            NUM_DIMENSIONS = getNumSVs(s2)
            svacts2 = np.dot(s2[:NUM_DIMENSIONS]*np.eye(NUM_DIMENSIONS), V2[:NUM_DIMENSIONS])

            if svacts1.shape != svacts2.shape:
                print("SVACTS1:", svacts1.shape)
                print("SVACTS2:", svacts2.shape)
            svcca_results = cca_core.get_cca_similarity(svacts1, svacts2, epsilon=1e-10, verbose=False)
            distances[i, j] = 1 - np.mean(svcca_results["cca_coef1"])
            distances[j, i] = distances[i, j]
            # distances are shape (numRNNs, numRNNs)

    # save the distances
    np.save("models/" + args.fname[:-4] + "_distances", distances)
else:
    distances = np.load("models/" + args.fname[:-4] + "_distances.npy")
plt.imshow(distances)
plt.clim(0, 1.0/3)
plt.colorbar()
plt.title("SVCCA Distances Between Networks")

clustering_algorithm = MDS()
clustered_data = clustering_algorithm.fit_transform(distances)

plt.show()

def Method1(rnn_distances, labels, class_counts):
    '''
    Model Agnostic Method for measuring how clustered Learning Rules are
    1.) Find the mean for each learning rule
    2.) Compute distance of all points for a given learning rule to the mean
        for that learning rule (d_lr)
    3.) Compute the distance of all points from other learning rules to the
        mean for the current learning rule (d_other)
    4.) Compute the ratio of d_lr / d_other
    '''
    N_RNNS = len(rnn_distances)
    # compute the mean/centroid for each learning rule
    centroid_positions = {0: np.zeros((N_RNNS)), 1: np.zeros((N_RNNS)), 2: np.zeros((N_RNNS)), 3: np.zeros((N_RNNS))}
    for rnn_ix in range(N_RNNS):  # loop over all RNNs
        centroid_positions[labels[rnn_ix]] += rnn_distances[rnn_ix]
    # now we must normalize the means
    for class_key in centroid_positions:
        centroid_positions[class_key] /= class_counts[class_key]
    # centroid_positions now contains the centroid for each learning rule

    # loop through all RNNs and get distance from centroids
    ooc_distances = {0:[], 1:[], 2:[], 3:[]}
    within_class_distances = {0:[], 1:[], 2:[], 3:[]}
    for rnn_ix in range(N_RNNS):
        # compute within class distance
        for class_key in centroid_positions:
            if class_key == labels[rnn_ix]:
                # compute within class distance\
                within_class_distances[class_key].append(ComputeDistance(centroid_positions[class_key], rnn_distances[rnn_ix]))
            else:
                # compute out of class distance
                ooc_distances[class_key].append(ComputeDistance(centroid_positions[class_key], rnn_distances[rnn_ix]))

    # get the mean of each list in both dictionaries and get the ratio
    ratios = []
    for class_key in within_class_distances:
        ratios.append(np.mean(within_class_distances[class_key])/np.mean(ooc_distances[class_key]))
    return ratios

def Method2(rnn_distances, labels, class_counts):
    '''
    Model Agnostic Method for measuring how clustered Learning Rules are
    1.) Compute average distance for any two points in the same learning rule
        (d_same)
    2.) Compute average distance between any two points (d_all)
    3.) Ratio of d_same / d_all
    '''
    N_RNNS = len(rnn_distances)

    # loop through all RNNs and get distance from all other RNNs
    ooc_distances = {0:[], 1:[], 2:[], 3:[]}
    within_class_distances = {0:[], 1:[], 2:[], 3:[]}
    for rnn_i in range(N_RNNS):
        for rnn_j in range(rnn_i+1, N_RNNS):
            # compute distance between RNN_i and RNN_j
            d = ComputeDistance(rnn_distances[rnn_i], rnn_distances[rnn_j])
            if labels[rnn_i] == labels[rnn_j]:  # update within class
                within_class_distances[labels[rnn_i]].append(d)
            else:  # update out-of-class
                ooc_distances[labels[rnn_i]].append(d)
                ooc_distances[labels[rnn_j]].append(d)
    # get the mean of each list in both dictionaries and get the ratio
    ratios = []
    for class_key in within_class_distances:
        ratios.append(np.mean(within_class_distances[class_key])/np.mean(ooc_distances[class_key]))
    return ratios

def Purity(rnn_distances, predicted, labels):
    '''
    Computes purity of clusters found with SVCCA

    Parameters
    ----------
    rnn_distances : obj
        model will be saved as a .pt file with this name in the /models/ directory.

    predicted : list<int>
        list of predicted class for each rnn

    labels : list<int>
        list of true class for each rnn

    Returns
    -------
    None.

    for each cluster,
        count the number of data points from the most common learning rule
        sum this accross all clusters
    divide the running sum by the total number of data points

    '''
    N_CLASSES = 4  # total number of rnn classes
    N_CLUSTERS = 4 # total number of rnn clusters (should be same as classes)
    N_RNNS = len(distances)  # total number of rnns clustered
    # dictionary counts the number of points belonging to each class
    # for each cluster. For example, class_counts[0][1] = 3 would
    # indicate that there were three points belonging to class 1 in
    # cluster 3. More generally,
    #    class_counts[cluster_id][class_id] = num_points_in_class_id_and_cluster_id
    #
    class_counts = {0:{0:0, 1:0, 2:0, 3:0},
                    1:{0:0, 1:0, 2:0, 3:0},
                    2:{0:0, 1:0, 2:0, 3:0},
                    3:{0:0, 1:0, 2:0, 3:0}}
    for ix in range(N_RNNS):  # loop through each RNN
        cluster_id = predicted[ix]  # cluster predicted by model
        class_id = labels[ix]       # true class of rnn
        # increment number of this class contained in this cluster
        class_counts[cluster_id][class_id] += 1

    # now we must sum the total number of points for each cluster belonging to the maximum
    # represented class
    max_per_class = np.zeros((N_CLUSTERS, 1))
    for cluster_id in range(N_CLASSES):
        # get the max in this dictionary
        max_pts_found = -1
        for key in class_counts[cluster_id]:
            if class_counts[cluster_id][key] > max_pts_found:  # current class contains the most rnns
                max_pts_found = class_counts[cluster_id][key]
        max_per_class[cluster_id] = max_pts_found
    # max_per_class now contains the number of points for the most represented class in
    # each cluster
    purity = np.sum(max_per_class, axis=0) / N_RNNS
    return purity

def getTrueLabelsHelper(numModelsOfType):
    '''
    numModelsOfType is a dictionary where
    each key is a model type and each value
    is the number of models of that type
    '''
    labels = []  # holds the ground truth labels
    class_counts = [0 for i in range(len(numModelsOfType)-1)]
    for class_id, key in enumerate(numModelsOfType):
        if key=="total":
            continue   # don't use totals
        curr_class_count = 0
        for j in range(numModelsOfType[key]):
            labels.append(class_id)
            curr_class_count += 1  # number of rnns in this class
        class_counts[class_id] = curr_class_count
    # note the labels in class_id don't have to
    # correspond with the labels used by the clustering
    # algorithm since we only consider purity and not
    # that the labels match
    return labels, class_counts

labels, class_counts = getTrueLabelsHelper(numModelsOfType)

if len(numModelsOfType) != 5 and len(numModelsOfType) != 3:
    N_colors = ['r', 'orange', 'g', 'b', 'indigo', 'violet']  # use when plotting against noise levels
else:
    N_colors = ['blue', 'orange', 'green', 'red']              # use when plotting the 4 learning rules
plt.figure(11)  # Scatter plot
count = 0     # keeps track of how many models plotted so far
ci = 0
legends = []  # holds model names to be used in legend
for modelType in numModelsOfType:
    if modelType == "total": # no plotting for total since it isn't a valid model type
        continue
    value = numModelsOfType[modelType]
    if modelType[0] == 'B':

        plt.scatter(clustered_data[count:count+value,0], clustered_data[count:count+value,1], c=N_colors[ci])
    else:

        plt.scatter(clustered_data[count:count+value,0], clustered_data[count:count+value,1], c=N_colors[ci])
    i+=1
    ci+=1
    count += value  # increment the number of models plotted
    legends.append(modelType) # add this model type name to the legend

plt.legend(legends)
plt.title("Sillhouette Score:" + str(silhouette_score(distances, labels)))

plt.figure()
x = np.array([ 1, 2, 3, 4, 5, 6])
y = np.array([0.25, 6.2e-4, 0.012, -0.013, -1.7e-5, -0.015])

m, b = np.polyfit(x, y, 1)
plt.plot(x, y, 'o')
plt.plot(x, m*x + b)
plt.ylim([-0.1, 0.5])


#!/usr/bin/env python
# coding: utf-8

# In[75]:

import os
import math
import operator
import numpy as np
import random
import pandas as pd
import GLAD
import glob, operator
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

# In[2]:

music = pd.read_csv('music_mturk_answers.csv')
sentiment = pd.read_csv('sentiment_mturk_answers.csv')
rte = pd.read_csv('rte.anno.csv')
data1 = music.values
data2 = sentiment.values
data3 = rte.values

# In[6]:

#for sentiment data(data2)
workers = np.unique(data2[:,0])
samples = np.unique(data2[:,1])

labelM = np.full((workers.shape[0],samples.shape[0]),-1)
for i in range(data2.shape[0]):
    x = int(np.where(workers == data2[i,0])[0])
    y = int(np.where(samples == data2[i,1])[0])
    labelM[x][y] = toint(data2[i][5])

# In[7]:

def toint(char):
    if char == 'pos':
        return 1
    else:
        return 0

# In[31]:

def getP(alpha,beta):
    pM = np.empty([])
    for i in range(alpha.shape[0]):
        for j in range(beta.shape[0]):
            pM = np.append(pM, 1/(1+math.exp(-alpha[i]*beta[j])))
            print(i)
    pM = pM.reshape(alpha.shape[0],beta.shape[0])
    return pM

# In[74]:

fig = plt.figure()
ax = plt.axes()
ax.set(xlabel='number of labels', ylabel='accuracy',title='sentiment')
x = [5,10,20,40,80,150,200]
y = [0.777,0.795,0.810,0.807,0.815,0.825,0.832]
ax.plot(x, y,'-g', label ='sentiment',color='blue')

# In[62]:

np.unique(data3[:,1]).shape

# In[93]:

labels = ['voting','spam','volcanoes']
id3 = [0.985,0.648,0.743]
lr = [0.959,0.605,0.635]
nb = [0.94,0.635,0.61]
glad = [0.988,0.689,0.765]

x = np.arange(len(labels))  # the label locations
width = 0.35

ig, ax = plt.subplots()
rects1 = ax.bar(x - width/2, id3, width/3, label='ID3')
rects2 = ax.bar(x - width/6, lr, width/3, label='LogR')
rects3 = ax.bar(x + width/6, nb, width/3, label='NaiveB', color = 'purple')
rects4 = ax.bar(x + width/2, glad, width/3, label='GLAD', color = 'red')

ax.set_ylabel('Accuracy')
ax.set_title('Accurcy comparesion')
ax.set_xticks(x)
ax.set_xticklabels(labels)
ax.legend()

def autolabel(rects):
    """Attach a text label above each bar in *rects*, displaying its height."""
    for rect in rects:
        height = rect.get_height()
        ax.annotate('{}'.format(height),
                    xy=(rect.get_x() + rect.get_width() / 2, height),
                    xytext=(0, 3),  # 3 points vertical offset
                    textcoords="offset points",
                    ha='center', va='bottom')

autolabel(rects1)
autolabel(rects2)

fig.tight_layout()

plt.show()

# In[ ]:


from gurobipy import GRB, quicksum, Model

# Initialise empty data
Nodes = []
Arcs = []
Omega = None
FromNode = -1
EPS = 0.00001
FlowINF = 10000

# Read in the data
with open("Outmax_flow2.dat") as f:
    while True:
        l = f.readline().strip()
        if l[0] == "n":
            FromNode = len(Nodes)
            Nodes.append(l)
        elif l[0] == "a":
            t = l.split(" ")
            flow = int(t[4])
            ToNode = int(t[3])
            if flow == FlowINF:
                Omega = len(Arcs)
            Arcs.append((FromNode, ToNode, flow))
        else:
            break
N = range(len(Nodes))
A = range(len(Arcs))

print("Nodes and Arcs", len(Nodes), len(Arcs))

# Empty array of jobs
Jobs = []
with open("max_flow2.dat_Job0") as f:
    while True:
        l = f.readline().strip()
        if len(l) <= 1:
            break
        nArc = int(l.split(" ")[1])
        t = f.readline().strip().split(" ")
        Jobs.append((nArc, int(t[1]) - 1, int(t[2]) - 1, int(t[3])))

T = range(1000)
J = range(len(Jobs))
# The jobs for an arc
JobsA = [[j for j in J if Jobs[j][0] == a] for a in A]
# The starting time periods for job A
JP = [range(Jobs[j][1], Jobs[j][2]) for j in J]

print("Jobs", len(Jobs))
# The jobs and their indices that impact on arc a in time t
JobsTA = [
    [
        {(j, p) for j in JobsA[a] for p in JP[j] if p <= t and p + Jobs[j][3] > t}
        for a in A
    ]
    for t in T
]
print("Calculated JobsTA")

# Set up the sub-problem
BSP = Model()
BSP.setParam("OutputFlag", 0)
X = {a: BSP.addVar() for a in A}

BSP.setObjective(X[Omega], GRB.MAXIMIZE)

FlowConservation = {
    n: BSP.addConstr(
        quicksum(X[a] for a in A if Arcs[a][0] == n)
        == quicksum(X[a] for a in A if Arcs[a][1] == n)
    )
    for n in N
}

UpperBound = {a: BSP.addConstr(X[a] <= Arcs[a][2]) for a in A}
# Solve unrestricted sub-problem to calculate MaxFlow
BSP.optimize()
MaxFlow = BSP.objVal
print("Maximum flow", MaxFlow)

# Set up master problem
BMP = Model()
Y = {(j, t): BMP.addVar(vtype=GRB.BINARY) for j in J for t in JP[j]}
Theta = {t: BMP.addVar(ub=MaxFlow) for t in T}

BMP.setObjective(quicksum(Theta[t] for t in T), GRB.MAXIMIZE)

EachJobScheduled = {j: BMP.addConstr(quicksum(Y[j, t] for t in JP[j]) == 1) for j in J}

_SolveBSP = {}
SolvesCalled = 0

def SolveBSP(AZero):
    global SolvesCalled
    SolvesCalled += 1
    if AZero not in _SolveBSP:
        for a in A:
            if a in AZero:
                UpperBound[a].RHS = 0
            else:
                UpperBound[a].RHS = Arcs[a][2]
        BSP.optimize()
        _SolveBSP[AZero] = (BSP.objVal, {a: UpperBound[a].pi for a in A})
    return _SolveBSP[AZero]

BMP._Best = 0

def Callback(model, where):
    if where == GRB.Callback.MIPSOL:
        YV = model.cbGetSolution(Y)
        YSet = {k for k in YV if YV[k] > 0.9}
        ThetaV = model.cbGetSolution(Theta)

        TotalObj = 0
        CutsAdded = 0
        for t in T:
            AZero = frozenset(a for a in A if JobsTA[t][a] & YSet)
            Obj, Pi = SolveBSP(AZero)
            TotalObj += Obj
            if Obj < ThetaV[t] - EPS:
                model.cbLazy(
                    Theta[t]
                    <= quicksum(
                        Pi[a]
                        * Arcs[a][2]
                        * (1 - quicksum(Y[j, tt] for (j, tt) in JobsTA[t][a]))
                        for a in A
                    )
                )
                CutsAdded += 1
            ThetaV[t] = Obj
        # Test if TotalObj is better than the incumbent
        if CutsAdded > 0 and TotalObj > model._Best + EPS:
            # Save the solution away
            print("#### Saving", TotalObj)
            model._Best = TotalObj
            model._Y = [YV[k] for k in Y]
            model._Theta = [ThetaV[k] for k in Theta]
        elif TotalObj > model._Best:
            model._Best = TotalObj

    if (
        where == GRB.Callback.MIPNODE
        and model.cbGet(GRB.Callback.MIPNODE_STATUS) == GRB.OPTIMAL
    ):
        # Do we have a saved solution that is better than the incumbent?
        if model._Best > model.cbGet(GRB.Callback.MIPNODE_OBJBST) + EPS:
            # If so, suggest it to Gurobi
            model.cbSetSolution(list(Y.values()), model._Y)
            model.cbSetSolution(list(Theta.values()), model._Theta)

# Warm start!
# Count = 1
# while True:
#     BSP.optimize()
#     if BSP.objVal>=2*MaxFlow:
#         break
#     print(Count, BSP.objVal)
#     Count+=1
#     for t in T:
#         BMP.addConstr(Theta[t]<=quicksum(
#                 UpperBound[a].pi*Arcs[a][2]*
#                 (1-quicksum(Y[j,tt] for (j,tt) in JobsTA[t][a]))
#                     for a in A))
#     for a in A:
#         if UpperBound[a].pi > EPS:
#             UpperBound[a].RHS += FlowINF

# Relax the problem
for k in Y:
    Y[k].vType = GRB.CONTINUOUS

# Add Benders cuts in a loop
BMP.setParam("OutputFlag", 0)
while True:
    BMP.optimize()
    CutsAdded = 0
    for t in T:
        for a in A:
            UpperBound[a].RHS = Arcs[a][2] * (
                1 - sum(Y[j, tt].x for (j, tt) in JobsTA[t][a])
            )
        BSP.optimize()
        if BSP.objVal < Theta[t].x - EPS:
            CutsAdded += 1
            BMP.addConstr(
                Theta[t]
                <= quicksum(
                    UpperBound[a].pi
                    * Arcs[a][2]
                    * (1 - quicksum(Y[j, tt] for (j, tt) in JobsTA[t][a]))
                    for a in A
                )
            )

    if CutsAdded == 0:
        break
    print("Added", CutsAdded, BMP.objVal)

BMP.setParam("OutputFlag", 1)
# Make the problem integer again
for k in Y:
    Y[k].vType = GRB.BINARY

BMP.setParam("MIPGap", 0)
BMP.setParam("LazyConstraints", 1)
# BMP.setParam('BranchDir',1)
BMP.optimize(Callback)


import sys
import json
from pathlib import Path
from bisect import bisect
import itertools
import collections

import lingpy
import segments
from pylexirumah import get_dataset

tokenizer = segments.tokenizer.Tokenizer()
model=lingpy.data.model.Model("asjp")
def to_asjp(segments):
    return lingpy.tokens2class(segments, model, cldf=False)
def convert_to_asjp(form):
    tokens = tokenizer(form, ipa=True).split()
    return to_asjp(tokens)
scorer = collections.defaultdict(lambda: -1)
scorer.update({(c, c): 0 for c in set(model.converter.values())})
def relative_edit_distance(x, y):
    a, b, s = lingpy.align.pairwise.pw_align(x, y, scale=1, scorer=scorer)
    return -s / len(a)

try:
    stats = json.load((Path(__file__).parent / "dataset_properties.json").open())
except (FileNotFoundError, json.decoder.JSONDecodeError):
    stats = {}

features = ["min_intersection", "average_intersection", "average_asjp_edit_distance", "mean_word_length", "mean_synonyms", "quantile_for_two_segments", "quantile_for_three_segments", "quantile_for_five_segments", "quantile_for_nine_segments", "n_lects", "mean_segments", "mean_asjp_segments", "mean_asjp_segments_product"]
features.sort()

def mean(x):
    return sum(x) / len(x)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--force", action='store_true')
    parser.add_argument("dataset", nargs="+")
    args = parser.parse_args()
    for argument in args.dataset:
        if argument in stats and not args.force:
            continue
        datafile = Path(argument)
        print(datafile)
        basename = datafile.stem

        dataset = get_dataset(datafile)
        c_language = dataset["FormTable", "languageReference"].name
        c_concept = dataset["FormTable", "parameterReference"].name
        c_form = dataset["FormTable", "form"].name
        c_segments = dataset["FormTable", "segments"].name
        lects = list(set(row[c_language] for row in dataset["FormTable"].iterdicts()))

        concepts = []
        wordlengths = []
        synonyms = []
        raw_segments = {}
        for l, language in enumerate(lects):
            print(language)
            c = {}
            concepts.append(c)
            for row in dataset["FormTable"].iterdicts():
                if row[c_language] == language:
                    c.setdefault(row[c_concept], []).append(row[c_form])
                    raw_segments.setdefault(language, set()).update(row[c_segments])
                    wordlengths.append(len(row[c_segments]))
            synonyms.extend([len(f) for f in c.values()])
        wordlengths.sort()
        asjp_segments = {language: set(to_asjp(list(phonemes)))
                         for language, phonemes in raw_segments.items()}

        intersections = []
        relative_edit_distances = []
        for concepts_i, concepts_j in itertools.combinations(concepts, 2):
            intersection = set(concepts_i).intersection(concepts_j)
            intersections.append(len(intersection))
            relative_concept_edit_distances = []
            for concept in intersection:
                for i, j in itertools.product(concepts_i[concept],
                                            concepts_j[concept]):
                    i = convert_to_asjp(i)
                    j = convert_to_asjp(j)
                    relative_concept_edit_distances.append(relative_edit_distance(i, j))
            if intersection:
                relative_edit_distances.append(
                    sum(relative_concept_edit_distances)/
                    len(relative_concept_edit_distances))

        print(basename)
        stats[argument] = {
            "min_intersection": min(intersections),
            "average_intersection": 2 * mean(intersections),
            "average_asjp_edit_distance": mean(relative_edit_distances),
            "mean_word_length": mean(wordlengths),
            "mean_synonyms": mean(synonyms),
            "quantile_for_two_segments": bisect(wordlengths, 2.01) / len(wordlengths),
            "quantile_for_three_segments": bisect(wordlengths, 3.01) / len(wordlengths),
            # In LexiRumah, a morpheme has 5.25 segments in the mean, and 90% of all marked morphemes have 9 segments or fewer.
            "quantile_for_five_segments": bisect(wordlengths, 5.25) / len(wordlengths),
            "quantile_for_nine_segments": bisect(wordlengths, 9.01) / len(wordlengths),
            "n_lects": len(lects),
            "mean_segments": mean([len(x) for x in raw_segments.values()]),
            "mean_asjp_segments": mean([len(x) for x in asjp_segments.values()]),
            "mean_asjp_segments_product": mean([
                len(x) * len(y) for x, y in itertools.combinations(asjp_segments.values(), 2)]),
        }

    json.dump(stats, Path(Path(__file__).parent / "dataset_properties.json").open("w"), indent=2, sort_keys=True)


from aiogram import Dispatcher, types
from aiogram.dispatcher import FSMContext
from aiogram.dispatcher.filters.state import State, StatesGroup
from aiogram.types import InlineKeyboardMarkup, InlineKeyboardButton, ReplyKeyboardRemove

from datetime import datetime
import logging

from database import memory_bot
from other import my_func

class FSMCommon(StatesGroup):
    authorization = State()
    authorization_enter_family = State()
    authorization_select_num = State()
    choice = State()

async def cmd_start(message: types.Message, state: FSMContext):
    await state.reset_state(with_data=True)
    async with state.proxy() as data:
        data["id_telegram"] = message.from_user.id
        kb = InlineKeyboardMarkup(row_width=1, resize_keyboard=True, one_time_keyboard=True)
        kb.insert(InlineKeyboardButton(text="Авторизоваться", callback_data="authorization"))
        await message.answer("Привет, {USERNAME}! Я буду помогать тебе с обучением.\n"
                             "Но, сначала мне нужно найти тебя в базе.\n"
                             "Нажми на кнопку ниже"
                             .format(USERNAME=message.from_user.first_name),
                             reply_markup=kb)
        await FSMCommon.authorization.set()

    await memory_bot.save_memory(state)

async def cmd_del_message(message: types.Message):
    await message.delete()

async def query_authorization(query: types.CallbackQuery, state: FSMContext):
    await query.message.delete()
    await query.message.answer("Напишите вашу ФАМИЛИЮ:", reply_markup=ReplyKeyboardRemove())
    await query.answer()
    await FSMCommon.authorization_enter_family.set()
    await memory_bot.save_memory(state)

async def cmd_authorization(message: types.Message, state: FSMContext):
    await state.reset_state(with_data=True)
    async with state.proxy() as data:
        data["id_telegram"] = message.from_user.id
    await message.answer("Напишите вашу ФАМИЛИЮ:", reply_markup=ReplyKeyboardRemove())
    await FSMCommon.authorization_enter_family.set()
    await memory_bot.save_memory(state)

async def authorization_enter_family(message: types.Message, state: FSMContext):
    user_data = await my_func.data_users_from_1c(message.text)
    if len(user_data) >= 1:
        text = "0: Меня нет в этом списке списке\n"
        for num, elem in enumerate(user_data):
            text += f"{num + 1}:{elem[0]}\n"
        await message.answer(f"{text}", reply_markup=types.ReplyKeyboardRemove())
        await message.answer(f"Вы под каким номером?")
        async with state.proxy() as data:
            data["User"] = user_data
        await FSMCommon.authorization_select_num.set()
        await memory_bot.save_memory(state)

    else:
        await message.answer(f"Что-то таких в списках нет, попробуйте написать еще раз или проверьте данные в 1с",
                             reply_markup=types.ReplyKeyboardRemove())

async def authorization_select_num(message: types.Message, state: FSMContext):
    if message.text.isdigit():
        async with state.proxy() as data:
            if 1 <= int(message.text) <= len(data["User"]):
                await message.answer(f"Отлично! {data['User'][int(message.text)-1][0]} так и запишем...")
                data["User_filial"] = data['User'][int(message.text)-1][2]
                data['User_otdel'] = data['User'][int(message.text)-1][1]
                data['User'] = data['User'][int(message.text) - 1][0]
            elif int(message.text) == 0:
                await message.answer("Вы точно написали вашу фамилию так-же как в 1С?\n"
                                     "Попробуйте написать еще раз.\nЕсли все равно не находит совпадения,"
                                     "то возможно вашу карточку еще не завели в сотрудниках 1С или допущена ошибка"
                                     , reply_markup=types.ReplyKeyboardRemove())
            else:
                await message.reply(f"Издеваетесь?) Выберите значение от 0 до {len(data['User'])}",
                                    reply_markup=types.ReplyKeyboardRemove())
                return

        await FSMCommon.choice.set()
        await memory_bot.save_memory(state)
        await variants_actions(message=message, state=state)
    else:
        await message.answer(f"Извините, давайте только по делу, укажите цифры с вашем ФИО:",
                             reply_markup=types.ReplyKeyboardRemove())

async def about_me(message: types.Message, state: FSMContext):
    await message.delete()
    async with state.proxy() as data:
        await message.answer("ID телеграмм: {}\n"
                             "ФИО: {}\n"
                             "Отдел: {}\n"
                             "Подразделение: {}".format(data.get("id_telegram"),
                                                        data.get("User"),
                                                        data.get("User_otdel"),
                                                        data.get("User_filial")),
                             reply_markup=types.ReplyKeyboardRemove())
        data["id_telegram"] = message.from_user.id

async def variants_actions(message: types.Message, state: FSMContext):
    async with state.proxy() as data:
        if not isinstance(data.get("User"), str):
            await message.delete()
            return
        data.pop("Current proc")
    kb = InlineKeyboardMarkup(row_width=2, resize_keyboard=True, one_time_keyboard=True)
    kb.insert(InlineKeyboardButton(text="Начать обучение", callback_data="start_teach"))
    kb.insert(InlineKeyboardButton(text="Начать наблюдение", callback_data="start_watch"))
    await message.answer("Что хочешь сделать?"
                         .format(USERNAME=message.from_user.first_name),
                         reply_markup=kb)
    await FSMCommon.choice.set()
    await memory_bot.save_memory(state)

async def query_variants_actions(query: types.CallbackQuery, state: FSMContext):
    if query.data == "start_teach":
        await query.message.answer("Напишите ФАМИЛИЮ обучаемого:", reply_markup=types.ReplyKeyboardRemove())
        async with state.proxy() as data:
            data["Current proc"] = {"Type": "Обучение", "Users": dict()}
        await state.set_state("FSMStudies:enter_family")
    elif query.data == "start_watch":
        await query.message.answer("Напишите ФАМИЛИЮ наблюдаемого:", reply_markup=types.ReplyKeyboardRemove())
        async with state.proxy() as data:
            data["Current proc"] = {"Type": "Наблюдение",
                                    "Users": dict(),
                                    "Time start": datetime.now().strftime("%d.%m.%Y %H:%M:%S"),
                                    "Time finish": None,
                                    "Duration": None,
                                    "message id to delete": [],
                                    "Check list": {
                                        "Установление контакта": {
                                            "1;Проявил инициативу и сам подошел к клиенту": [0, 2],
                                            "2;Поприветствовал покупателя, представился": [0, 1, 2],
                                            "3;Применил фразу на установление контакта"
                                            " Клиент охотно вступил в диалог": [0, 2],
                                            "4;Узнал имя покупателя": [0, 2],
                                        },
                                        "Выяснение потребностей": {
                                            "5;Уточнил запрос клиента, используя воронку вопросов. "
                                            "опросы построены грамотно и последовательно": [0, 1, 3],
                                            "6;Резюмировал выявленные потребности": [0, 3],
                                        },
                                        "Презентация": {
                                            "7;Сотрудник хорошо ориентируется в ассортименте": [0, 3],
                                            "8;Говорит простым языком, объясняет терминологию": [0, 1],
                                            "9;В предложении товара были использованы выявленные точки воздействия":
                                                [0, 1, 3],
                                            "10;Сделал комплексное предложение (как минимум, два "
                                            "дополнительных товара, которые могут подойти)": [0, 2],
                                        },
                                        "Работа с возражениями": {
                                            "11;Работал с возражением по алгоритму": [0, 2],
                                            "12;Говорил убедительно и грамотно, опираясь на потребности клиента":
                                                [0, 1, 2],
                                        },
                                        "Завершение диалога": {
                                            "13;Предложил дополнительные услуги "
                                            "(нарезка, монтаж, доставка, подъем на этаж)": [0, 2],
                                            "14;Договорился с клиентом о дальнейших действиях. Пригласил приехать в"
                                            " магазин или договарился о следующем звонке": [0, 2],
                                            "15;Вежливо завершает разговор": [0, 1],
                                        },
                                        "Общее впечатление ": {
                                            "16;Доброжелательное выражение лица и интонация на "
                                            "протяжении всего диалога": [0, 2],
                                            "17;Нет долгих пауз и неуместных повторений вопросов ": [0, 2],
                                            "18;Производит впечатление профессионала, убедителен в коммуникации": [0, 3],
                                        }
                                    }}
        await state.set_state("FSMObservation:enter_family")
    else:
        return
    await query.message.delete()
    await memory_bot.save_memory(state)

def register_handlers(dp: Dispatcher):
    dp.register_message_handler(cmd_start, commands="start", state="*")
    dp.register_message_handler(cmd_authorization, commands="authorization", state="*")
    dp.register_message_handler(about_me, commands="about_me", state="*")
    dp.register_message_handler(variants_actions, commands="actions", state="*")
    dp.register_callback_query_handler(query_variants_actions, state=FSMCommon.choice)

    dp.register_message_handler(cmd_del_message, state=[FSMCommon.authorization, FSMCommon.choice,
                                                        "FSMStudies:add_user_or_next_step",
                                                        "FSMStudies:choice_train",
                                                        "FSMStudies:resume",
                                                        "FSMStudies:time"])
    dp.register_callback_query_handler(query_authorization, state=FSMCommon.authorization)
    dp.register_message_handler(authorization_enter_family, state=FSMCommon.authorization_enter_family)
    dp.register_message_handler(authorization_select_num, state=FSMCommon.authorization_select_num)


import torch
import torch.nn as nn
from mover.utils.camera import (
    get_pitch_roll_euler_rotation_matrix,
)

### ground plane nformation
def set_ground_plane(self, ground_plane):
    self.ground_plane.data.copy_(ground_plane)

# def get_ground_plane(self):
#     verts = self.get_verts_object_parallel_ground()
#     return torch.max(verts[:, :, 1])

def get_ground_plane_mesh(self):
    gp_y = self.ground_plane
    vertices = torch.Tensor([[[-3, 0, 0], [3, 0, 0], [3, 0, 7], [-3, 0, 7]]]).type_as(gp_y)
    vertices[:, :, 1] = gp_y
    faces = torch.Tensor([[[0, 1, 2], [0, 2, 3]]]).type_as(self.faces)
    green = [0.19607843, 0.80392157, 0.19607843]
    textures = torch.Tensor([green, green]).reshape(1, 2, 1, 1, 1, 3).type_as(self.textures)
    return vertices, faces, textures

def get_ground_plane_np(self):
    vertices, faces, textures = self.get_ground_plane_mesh()
    return vertices.detach().cpu().numpy(), faces.detach().cpu().numpy(), textures.detach().cpu().numpy()

## camera -> image
def apply_projection_to_image(self, verts):
    proj_v = torch.matmul(self.K_intrin, verts.transpose(2,1)).transpose(2,1)
    xy = proj_v[:, :, :2] / proj_v[:, :, 2:]
    z = proj_v[:, :, 2:]
    return xy, z

def apply_projection_to_local_image(self, verts, obj_cam=None, resize=None):
    assert obj_cam is not None
    proj_v = torch.matmul(self.K_rois[obj_cam].float(), verts.float().transpose(2,1)).transpose(2,1)
    xy = proj_v[:, :, :2] / proj_v[:, :, 2:]
    z = proj_v[:, :, 2:]

    return xy, z

def get_relative_extrin_new(self):
    return get_pitch_roll_euler_rotation_matrix(self.rotate_cam_pitch, self.rotate_cam_roll) # need in forward

# camera
def get_cam_extrin(self):
    extra_extrin_mat = self.get_relative_extrin_new()
    cam_extrin = torch.transpose(torch.matmul(extra_extrin_mat, torch.transpose(self.K_extrin, 2, 1)), 2, 1)
    return cam_extrin


# From baselines.bench.monitor.py. Here so the client outside docker does not need to install baselines.
# Only pandas, tkinter (sudo apt-get install python3-tk), numpy, matplotlib.
# cd ~/custom_roboschool/servorobots/servorobots/tools
# python3 plot_results.py --dirs /progress_1
from glob import glob
import os.path as osp
import json

def load_results(dir):
    import pandas
    monitor_files = (
        glob(osp.join(dir, "*monitor.json")) +
        glob(osp.join(dir, "*monitor.csv"))) # get both csv and (old) json files
    if not monitor_files:
        raise LoadMonitorResultsError("no monitor files of the form *%s found in %s" % (Monitor.EXT, dir))
    dfs = []
    headers = []
    for fname in monitor_files:
        with open(fname, 'rt') as fh:
            if fname.endswith('csv'):
                firstline = fh.readline()
                if not firstline:
                    continue
                assert firstline[0] == '#'
                header = json.loads(firstline[1:])
                df = pandas.read_csv(fh, index_col=None)
                headers.append(header)
            elif fname.endswith('json'): # Deprecated json format
                episodes = []
                lines = fh.readlines()
                header = json.loads(lines[0])
                headers.append(header)
                for line in lines[1:]:
                    episode = json.loads(line)
                    episodes.append(episode)
                df = pandas.DataFrame(episodes)
            else:
                assert 0, 'unreachable'
            df['t'] += header['t_start']
        dfs.append(df)
    df = pandas.concat(dfs)
    df.sort_values('t', inplace=True)
    df.reset_index(inplace=True)
    df['t'] -= min(header['t_start'] for header in headers)
    df.headers = headers # HACK to preserve backwards compatibility
    return df

import numpy as np
import matplotlib
matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode

import matplotlib.pyplot as plt
plt.rcParams['svg.fonttype'] = 'none'

X_TIMESTEPS = 'timesteps'
X_EPISODES = 'episodes'
X_WALLTIME = 'walltime_hrs'
Y_REWARD = 'reward'
Y_TIMESTEPS = 'timesteps'
POSSIBLE_X_AXES = [X_TIMESTEPS, X_EPISODES, X_WALLTIME]
EPISODES_WINDOW = 90
COLORS = ['blue', 'green', 'red', 'cyan', 'magenta', 'yellow', 'black', 'purple', 'pink',
        'brown', 'orange', 'teal', 'coral', 'lightblue', 'lime', 'lavender', 'turquoise',
        'darkgreen', 'tan', 'salmon', 'gold', 'lightpurple', 'darkred', 'darkblue']

def rolling_window(a, window):
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

def window_func(x, y, window, func):
    yw = rolling_window(y, window)
    yw_func = func(yw, axis=-1)
    return x[window-1:], yw_func

def ts2xy(ts, xaxis, yaxis):
    if xaxis == X_TIMESTEPS:
        x = np.cumsum(ts.l.values)
    elif xaxis == X_EPISODES:
        x = np.arange(len(ts))
    elif xaxis == X_WALLTIME:
        x = ts.t.values / 3600.
    else:
        raise NotImplementedError
    if yaxis == Y_REWARD:
        y = ts.r.values
    elif yaxis == Y_TIMESTEPS:
        y = ts.l.values
    else:
        raise NotImplementedError
    return x, y

def plot_curves(fig, xy_list, xaxis, yaxis, title):

    maxx = max(xy[0][-1] for xy in xy_list)
    minx = 0
    for (i, (x, y)) in enumerate(xy_list):
        color = COLORS[i]
        plt.scatter(x, y, s=2)
        x, y_mean = window_func(x, y, EPISODES_WINDOW, np.mean) #So returns average of last EPISODE_WINDOW episodes
        plt.plot(x, y_mean, color=color)
    plt.xlim(minx, maxx)
    plt.title(title)
    plt.xlabel(xaxis)
    plt.ylabel(yaxis)
    plt.tight_layout()
    fig.canvas.mpl_connect('resize_event', lambda event: plt.tight_layout())
    plt.grid(True)

def plot_results(dirs, num_timesteps, xaxis, task_name):
    tslist = []
    for dir in dirs:
        ts = load_results(dir)
        ts = ts[ts.l.cumsum() <= num_timesteps]
        tslist.append(ts)
    yaxes = ['reward', 'timesteps']
    fig = plt.figure(figsize=(8, 4))
    for i in range(len(yaxes)):
        plt.subplot(2, 1, i+1)
        xy_list = [ts2xy(ts, xaxis, yaxes[i]) for ts in tslist]
        plot_curves(fig, xy_list, xaxis, yaxes[i], task_name)

# Example usage in jupyter-notebook
# from baselines import log_viewer
# %matplotlib inline
# log_viewer.plot_results(["./log"], 10e6, log_viewer.X_TIMESTEPS, "Breakout")
# Here ./log is a directory containing the monitor.csv files

def main():
    import argparse
    import os
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--dirs', help='List of log directories', nargs = '*', default=['./log'])
    parser.add_argument('--num_timesteps', type=int, default=int(10e6))
    parser.add_argument('--xaxis', help = 'Varible on X-axis', default = X_TIMESTEPS)
    parser.add_argument('--task_name', help = 'Title of plot', default = 'Breakout')
    parser.add_argument('--show', help="Show the window", default = True, dest='show', action='store_true')
    parser.add_argument('--no-show', help="Don't show the window", dest='show', action='store_false')
    parser.add_argument('--save_dir', help='Show the window')
    parser.add_argument('--format', help='png, pdf (+ what is supported by matplotlib)', default='png')
    args = parser.parse_args()
    args.dirs = [os.path.abspath(dir) for dir in args.dirs]
    plot_results(args.dirs, args.num_timesteps, args.xaxis, args.task_name)
    if args.show:
        print('showing plot')
        plt.show()
    if args.save_dir is not None:
        plt.savefig(args.save_dir + '/results.' + args.format)
        print('Saved plot in ' + args.save_dir + '/results.' + args.format)

if __name__ == '__main__':
    main()


import logging
import math
import time
import warnings
from itertools import cycle, islice

import torch
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision import transforms
from sklearn.preprocessing import StandardScaler
from sklearn.datasets import make_moons

from .augmentation import RandAugment, CutOutAug
from .autoaugment import CIFAR10Policy, SVHNPolicy
from .utils import AttrDict

cifar10_mean = (0.4914, 0.4822, 0.4465)
cifar10_std = (0.2471, 0.2435, 0.2616)
cifar100_mean = (0.5071, 0.4867, 0.4408)
cifar100_std = (0.2675, 0.2565, 0.2761)
normal_mean = (0.5, 0.5, 0.5)
normal_std = (0.5, 0.5, 0.5)

class MyDataset(Dataset):
    def __init__(self, data, targets, transform=None):
        self.data = data
        self.targets = torch.LongTensor(targets)
        self.transform = transform

    def __getitem__(self, index):
        x = self.data[index]
        y = self.targets[index]

        if self.transform:
            x = self.transform(x)
        return x, y

    def __len__(self):
        return len(self.data)

class TransformToy(object):
    def __init__(self):
        self.ori = transforms.Compose([])
        self.aug = transforms.Compose([])

    def __call__(self, x):
        ori = self.ori(x)
        aug = self.aug(x)
        return ori, aug

def prepare_cifar10(num_labeled=4000, classes=10, eval_step=1000, batch_size=64, resize=32, data_path="./data",randaug=[2,16]):
    args = {"num_labeled": num_labeled,
        "num_classes": classes,
        "eval_step": eval_step,
        "batch_size": batch_size,
        "resize": resize,
        "data_path": data_path,
        "expand_labels": False,
        "randaug": randaug
       }
    args = AttrDict(args)
    return get_cifar10(args)

def get_toy(args):
    n_samples = 1500
    np.random.seed(0)
    noisy_moons = make_moons(n_samples=n_samples, noise=0.05)
    X, y = noisy_moons

    X = StandardScaler().fit_transform(X)
    random_choice = np.random.choice(X.shape[0], 5, replace=False)
    labeled_data = X[random_choice]
    targets  = y[random_choice]

    labeled_dataset = MyDataset(labeled_data, torch.from_numpy(targets).long())
    unlabeled_dataset =MyDataset(X, torch.from_numpy(y).long(), transform=TransformToy())
    test_dataset =MyDataset(X, torch.from_numpy(y).long())

    return labeled_dataset, unlabeled_dataset, test_dataset

def get_cifar10(args):
    transform_labeled = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        CIFAR10Policy(),
        transforms.RandomCrop(size=args.resize,
                              padding=int(args.resize*0.125),
                              padding_mode='reflect'),
        transforms.ToTensor(),
        transforms.Normalize(mean=cifar10_mean, std=cifar10_std)
    ])
    transform_val = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=cifar10_mean, std=cifar10_std)
    ])
    base_dataset = datasets.CIFAR10(args.data_path, train=True, download=True)

    train_labeled_idxs, train_unlabeled_idxs = x_u_split(args, base_dataset.targets)

    train_labeled_dataset = CIFAR10SSL(
        args.data_path, train_labeled_idxs, train=True,
        transform=transform_labeled
    )

    train_unlabeled_dataset = CIFAR10SSL(
        args.data_path, train_unlabeled_idxs,
        train=True,
        transform=TransformMPL(args, mean=cifar10_mean, std=cifar10_std)
    )

    test_dataset = datasets.CIFAR10(args.data_path, train=False,
                                    transform=transform_val, download=False)

    return train_labeled_dataset, train_unlabeled_dataset, test_dataset

def get_cifar100(args):

    transform_labeled = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        CIFAR10Policy(),
        transforms.RandomCrop(size=args.resize,
                              padding=int(args.resize*0.125),
                              padding_mode='reflect'),
        transforms.ToTensor(),
        transforms.Normalize(mean=cifar100_mean, std=cifar100_std)])

    transform_val = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=cifar100_mean, std=cifar100_std)])

    base_dataset = datasets.CIFAR100(args.data_path, train=True, download=True)

    train_labeled_idxs, train_unlabeled_idxs = x_u_split(args, base_dataset.targets)

    train_labeled_dataset = CIFAR100SSL(
        args.data_path, train_labeled_idxs, train=True,
        transform=transform_labeled
    )

    train_unlabeled_dataset = CIFAR100SSL(
        args.data_path, train_unlabeled_idxs, train=True,
        transform=TransformMPL(args, mean=cifar100_mean, std=cifar100_std)
    )

    test_dataset = datasets.CIFAR100(args.data_path, train=False,
                                     transform=transform_val, download=False)

    return train_labeled_dataset, train_unlabeled_dataset, test_dataset

def x_u_split(args, labels):
    label_per_class = args.num_labeled // args.num_classes
    labels = np.array(labels)
    labeled_idx = []
    # unlabeled data: all training data
    unlabeled_idx = np.array(range(len(labels)))
    for i in range(args.num_classes):
        idx = np.where(labels == i)[0]
        idx = np.random.choice(idx, label_per_class, False)
        labeled_idx.extend(idx)
    labeled_idx = np.array(labeled_idx)
    assert len(labeled_idx) == args.num_labeled

    if args.expand_labels or args.num_labeled < args.batch_size:
        num_expand_x = math.ceil(
            args.batch_size * args.eval_step / args.num_labeled)
        labeled_idx = np.hstack([labeled_idx for _ in range(num_expand_x)])
    np.random.shuffle(labeled_idx)
    return labeled_idx, unlabeled_idx

class TransformMPL(object):
    def __init__(self, args, mean, std):
        if args.randaug:
            n, m = args.randaug
        else:
            n, m = 2, 16  # default

        self.ori = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=args.resize,
                                  padding=int(args.resize*0.125),
                                  padding_mode='reflect')])
        self.aug = transforms.Compose([
            CIFAR10Policy(),
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=args.resize,
                                  padding=int(args.resize*0.125),
                                  padding_mode='reflect'),
            RandAugment(n=n, m=m),
            CutOutAug(n=1, size=args.resize//4),
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(size=args.resize,
                                  padding=int(args.resize*0.125),
                                  padding_mode='reflect')
            ])
        self.normalize = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)])

    def __call__(self, x):
        ori = self.ori(x)
        aug = self.aug(x)
        return self.normalize(ori), self.normalize(aug)

class CIFAR10SSL(datasets.CIFAR10):
    def __init__(self, root, indexs, train=True,
                 transform=None, target_transform=None,
                 download=False):
        super().__init__(root, train=train,
                         transform=transform,
                         target_transform=target_transform,
                         download=download)
        if indexs is not None:
            self.data = self.data[indexs]
            self.targets = np.array(self.targets)[indexs]

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

class CIFAR100SSL(datasets.CIFAR100):
    def __init__(self, root, indexs, train=True,
                 transform=None, target_transform=None,
                 download=False):
        super().__init__(root, train=train,
                         transform=transform,
                         target_transform=target_transform,
                         download=download)
        if indexs is not None:
            self.data = self.data[indexs]
            self.targets = np.array(self.targets)[indexs]

    def __getitem__(self, index):
        img, target = self.data[index], self.targets[index]
        img = Image.fromarray(img)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target


from subprocess import call
from tkinter import filedialog
import subprocess
import requests
import os
import platform

class Application(object):

    app = None
    viewer = None
    smbUser = ('user', 'password')

    serverUrlWidget = None
    versionsWidget = None
    installPathWidget = None
    runWidget = None
    networkStatusWidget = None
    letterListWidget = None
    bootModeWidget = None
    archWidget = None

    winTarget = ''
    archTarget = ''
    bootMode = ''

    jsonConfig = 'netinstall.json'

    availablePaths = {}

    def __init__(self, smbUser, winTarget, createViewHandle, args):

        Application.app = self

        dScaling = 0
        isDBoot = False

        for i in args:

            if "dboot=" in i:
                isDBoot = i.split("=")[1] == "1"
                continue

            if "dscaling=" in i:
                dScaling = float(i.split("=")[1]);
                continue

        self.bootMode = self.getBootType(isDBoot)

        arch = platform.architecture()[0]
        if arch == '64bit':
            self.archTarget = 'x64'
        elif arch == '32bit':
            self.archTarget = 'x86'
        else:
            self.archTarget = 'Desconhecida'

        self.smbUser = smbUser
        self.winTarget = winTarget
        self.viewer = createViewHandle(dScaling)

    def init(self):

        self.serverUrlWidget = self.viewer.getWidgetById('serverUrl')
        self.versionsWidget = self.viewer.getWidgetById('version-list')
        self.installPathWidget = self.viewer.getWidgetById('install-path')
        self.runWidget = self.viewer.getWidgetById('run')
        self.networkStatusWidget = self.viewer.getWidgetById('network-status')
        self.letterListWidget = self.viewer.getWidgetById('letter-list')
        self.bootModeWidget = self.viewer.getWidgetById('boot-mode')
        self.archWidget = self.viewer.getWidgetById('arch')

        self.runWidget['state'] = 'disabled'
        self.viewer.addCLickListener('run', Application.onClick)
        self.viewer.addCLickListener('load-driver', Application.onClick)
        self.viewer.addCLickListener('run-diskpart', Application.onClick)
        self.viewer.addCLickListener('run-cmd', Application.onClick)

        self.versionsWidget.bind('<<ComboboxSelected>>', Application.onSelectVerion);
        self.viewer.addCLickListener('refresh', Application.onClick)

        letters = []
        for l in range(ord('A'), ord('Z') + 1):
            l = chr(l)+':'
            if not self.letterIsUsed(l):
                letters.append(l)

        self.letterListWidget['values'] = letters
        self.letterListWidget.set(self.letterListWidget['values'][0])

        self.bootModeWidget['text'] = self.bootMode
        self.archWidget['text'] = self.archTarget

    def run(self):
        self.refresh()
        self.viewer.loop()

    def onDetectNetWork(self):
        self.networkStatusWidget['text'] = 'Conectado'
        print('* Conexão detectada!')

    def refresh(self):

        self.availablePaths = {}
        self.versionsWidget['values'] = ()
        serverUri = self.serverUrlWidget.get()

        print('* Verificando ping com "'+serverUri+'"...')
        if Application.runProg(['ping', serverUri, '-n', '1' , '-w', '250'], False) != 0 :
            print('* Sem conxeão!')
            return

        print('\n* Baixando arquivo json em "http://'+serverUri+'/'+ self.jsonConfig +'"...')

        response = requests.get('http://' + serverUri + '/' + self.jsonConfig)
        if response.status_code != 200:
            print('* Erro ao baixar o arquivo "'+self.jsonConfig+'".')
            print('* Codigo http: ' + str(response.status_code) + '.')
            return

        cfg = response.json();
        if self.winTarget not in cfg:
            print('* Erro ao buscar a versão alvo para o windows "' + self.winTarget + '".')
            return

        self.onDetectNetWork()
        print('* Versões disponíveis:')

        listValues = ()
        for i in cfg[self.winTarget]:

            item = self.winTarget+'-'+i
            i = cfg[self.winTarget][i]

            if self.archTarget not in i:
                continue

            listValues = listValues + (item,)

            self.availablePaths[item] = {
                'path': i[self.archTarget]['path'],
                'exec': i[self.archTarget]['exec']
            }

            print('\t- ' + item)

        if len(listValues) > 1 :
            self.versionsWidget['values'] = listValues
            self.versionsWidget.set(listValues[0])
            Application.onSelectVerion(None)
        else:
            print('\t - Nenhuma versão disponível')

    def onRunBtnClick(self):

        letter = self.letterListWidget.get();
        serverUri = self.serverUrlWidget.get();
        path = self.installPathWidget.get();
        item = Application.app.versionsWidget.get()
        exe = Application.app.availablePaths[item]['exec']

        if self.letterIsUsed(letter):
            if not os.path.isfile(letter + '\\'+ exe):
                print('* A letra de unidade "' + letter + '" já está em uso. Tente uma outra.')
        else:
            if not self.mountShared(letter, serverUri, path, self.smbUser):
                return

        Application.runProg([letter + '\\' + exe], True)

    def mountShared(self, letter, serverUri, path, user):

        url = '\\\\'+serverUri+'\\'+path

        print('* Tentando montar "'+url+'" em '+letter+'...')
        code = Application.runProg(['net', 'use', letter, url, '/u:'+user[0], user[1]], False)
        if code != 0:
            print('* Erro ao tentar montar unidade de rede.')
            print('* Codigo de erro: ' + str(code) + '.')
            return False

        print('* Unidade montada com sucesso!')
        return True

    def letterIsUsed(self, letter):
        return os.path.isdir(letter)

    def onLoadDriverBtnClick(self):

        driverFile = filedialog.askopenfilename(title = "Selecione o arquivo de driver...", filetypes = (("Driver .inf","*.inf"), ("all files","*.*")))

        if not driverFile:
            return

        print('* Carregando arquivo de driver "'+driverFile+'"...')
        code = Application.runProg(['drvload', driverFile], False)

        if code != 0:
            print('* Erro ao carregar driver com código '+ str(code) +'.')
        else:
            print("* Driver carregado com sucesso!")
            self.refresh()

    @staticmethod
    def onClick(id, viewer, event):

        if id == 'refresh':
            Application.app.refresh()

        if id == 'run':
            Application.app.onRunBtnClick()

        if id == 'load-driver':
            Application.app.onLoadDriverBtnClick()

        if id == 'run-diskpart':
            Application.runProg(['start.bat', 'diskpart'], True)

        if id == 'run-cmd':
            Application.runProg(['start.bat', 'cmd'], True)

    @staticmethod
    def onSelectVerion(event):

        item = Application.app.versionsWidget.get()
        path = Application.app.availablePaths[item]['path']

        Application.app.installPathWidget['state'] = 'normal'
        Application.app.installPathWidget.delete(0, 'end')
        Application.app.installPathWidget.insert(0, path)
        Application.app.installPathWidget['state'] = 'readonly'

        Application.app.runWidget['state'] = 'normal'

    @staticmethod
    def runProg(entry, independent):
        if independent:
            return subprocess.Popen(entry)
        else:
            return call(entry)

    @staticmethod
    def getBootType(isDebug):

        if isDebug :
            return "DEBUG"

        output = subprocess.check_output(['boot-mode.bat'])
        if output == None :
            return 'Desconhecido'

        bootTypeCode = str(output).split(' ')[-1].split('\\r')[0];

        if bootTypeCode == '0x1':
            return 'BIOS'
        elif bootTypeCode == '0x2':
            return 'UEFI'
        else:
            return 'Desconhecido'


# -*- coding: utf-8 -*-
# Author : zenk
# 2013-07-05 20:32
from os import path
from collections import defaultdict
import re

from PyQt4.QtGui import QDialog
from PyQt4.QtGui import QFileDialog
from PyQt4.QtGui import QGridLayout
from PyQt4.QtGui import QLineEdit
from PyQt4.QtGui import QLabel
from PyQt4.QtGui import QIntValidator
from PyQt4.QtGui import QMessageBox
from PyQt4.QtGui import QPushButton
from PyQt4.QtGui import QVBoxLayout

from PyQt4.QtCore import SIGNAL

import xlrd

class ReplaceI18NDialog(QDialog):
  def __init__(self, config_name, parent=None):
    super(ReplaceI18NDialog, self).__init__(parent)
    self.parent = parent
    grid_layout = QGridLayout()
    grid_layout.addWidget(QLabel(u'表单名字/序号'), 0, 0)
    self.sheet_index_input = QLineEdit(config_name, self)
    grid_layout.addWidget(self.sheet_index_input, 0, 1)

    grid_layout.addWidget(QLabel(u'列'), 1, 0)
    self.column_input = QLineEdit("1", self)
    int_validator = QIntValidator()
    int_validator.setRange(1, 100)
    self.column_input.setValidator(int_validator)
    grid_layout.addWidget(self.column_input, 1, 1)

    grid_layout.addWidget(QLabel(u'文件'), 2, 0)
    self.file_input = QLineEdit("", self)
    grid_layout.addWidget(self.file_input, 2, 1)
    select_file_btn = QPushButton(u'选择')
    grid_layout.addWidget(select_file_btn, 2, 2)

    do_btn = QPushButton(u'替换')
    main_layout = QVBoxLayout()
    main_layout.addLayout(grid_layout)
    main_layout.addWidget(do_btn)
    self.setLayout(main_layout)

    self.connect(select_file_btn, SIGNAL("clicked()"), self.select_file)
    self.connect(do_btn, SIGNAL("clicked()"), self.replace)

  def select_file(self):
    filename = unicode(QFileDialog.getOpenFileName(self, "Open File",
                                           "", "Excel (*.xls)"))
    if not path.exists(filename):
      QMessageBox.information(self, u"i18n替换",
          u'文件:%s不存在' % filename)
      return
    book = xlrd.open_workbook(filename)
    sheet_name = unicode(self.sheet_index_input.text())
    sheet = None
    if sheet_name in book.sheet_names():
      sheet = book.sheet_by_name(sheet_name)
    elif re.match("\d+", sheet_name):
      sheet = book.sheet_by_index(int(sheet_name) - 1)
    if not sheet:
      QMessageBox.information(self, u'i18n替换',
          u'表单名字/索引:%s不存在' % sheet_name)
      return
    self.text_dic = defaultdict(str)
    for k, v in zip(sheet.col_values(0),
        sheet.col_values(int(unicode(self.column_input.text())))):
        self.text_dic[k] = v
    self.file_input.setText(filename)

  def replace(self):
    not_replaced = self.parent.replace_i18n(self.text_dic)
    QMessageBox.information(self, u'i18n替换',
        u'替换成功%s' % ('' if not not_replaced else u'\n以下字符串没有被替换:%s' % not_replaced))
    self.close()


from loader import dp
from aiogram.types import Message
import game
import random

@dp.message_handler()
async def mes_help(message: Message):
    if game.new_game:
        for duel in game.total:
            if message.from_user.id == duel[0]:
                count = message.text
                name = message.from_user.first_name
                if count.isdigit() and 0 < int(count) <29:
                    duel[2] -= int(count)
                    await check_win(message, 'Ты', duel)
                    if game.new_game:
                        await message.answer(f'{duel[1]} взял {count} конфет и на столе осталось {duel[2]}\n Теперь ход бота...')

                    bot_take = random.randint(1,28) if duel[2] >28 else duel[2]
                    duel[2] -= bot_take

                    await check_win(message, 'Бот', duel)
                    if game.new_game:
                        await message.answer(f'Бот взял {bot_take} конфет и на столе осталось {duel[2]}\n Теперь твой ход...')

            else:
                await message.answer(f'Введите число от 1 до 28')

async def check_win(message: Message, win: str, total: list):
    if total[2] <= 0:
        await message.answer(f'{win} Победил!')
        game.new_game = False
        game.total.remove(total)


import cv2
import numpy as np
import matplotlib.pyplot as plt
def region_of_interest(img, vertices):
    mask = np.zeros_like(img)
    # channel_count = img.shape[2]
    match_mask_color = 255
    cv2.fillPoly(mask, vertices, match_mask_color)
    masked_image = cv2.bitwise_and(img, mask)
    return masked_image
def draw_lines(img,lines):
    for line in lines:
        (x1,y1,x2,y2)=line[0]
        cv2.line(img,(x1,y1),(x2,y2),(255,0,255),6)
    return img
# img=cv2.imread('../test_road.jpg')
# img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
def process(img):
    print(img.shape)
    height=img.shape[0]
    width=img.shape[1]
    ROI=[(0,height),(width/2,height/2),(width,height)]

    gray_image=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    edge=cv2.Canny(gray_image,70,200)
    crop_image=region_of_interest(edge,np.array([ROI],np.int32))
    lines=cv2.HoughLinesP(crop_image,6,np.pi/60,160,np.array([]),minLineLength=40,maxLineGap=25)
    f_image=draw_lines(img,lines)
    return  f_image
cap=cv2.VideoCapture('../test1.mp4')

while (cap.isOpened()):
    _,frame=cap.read()
    frame=process(frame)
    cv2.imshow('test_road',frame)
    if cv2.waitKey(1) &0xFF ==ord('q'):
        break


from scipy.ndimage import zoom
from torchvision import transforms
from PIL import Image
from cv2 import cv2
import cv2
import numpy as np
import torch

def gaussian_filter(img, K_size=7, sigma=0.4):
    if len(img.shape) == 3:
        H, W, C = img.shape
    else:
        img = np.expand_dims(img, axis=-1)
        H, W, C = img.shape
    ## Zero padding
    pad = K_size // 2
    out = np.zeros((H + pad * 2, W + pad * 2, C), dtype=np.float)
    out[pad: pad + H, pad: pad + W] = img.copy()
    ## prepare Kernel
    K = np.zeros((K_size, K_size), dtype=np.float)
    for x in range(-pad, -pad + K_size):
        for y in range(-pad, -pad + K_size):
            K[y + pad, x + pad] = np.exp(-(x ** 2 + y ** 2) / (2 * (sigma ** 2)))
    K /= (2 * np.pi * sigma * sigma)
    K /= K.sum()
    tmp = out.copy()
    # filtering
    for y in range(H):
        for x in range(W):
            for c in range(C):
                out[pad + y, pad + x, c] = np.sum(K * tmp[y: y + K_size, x: x + K_size, c])
    out = np.clip(out, 0, 255)
    out = out[pad: pad + H, pad: pad + W]
    return out

def show_cam_on_image(img, mask):
    heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255
    cam = heatmap + np.float32(img)
    cam = cam / np.max(cam)
    cv2.imwrite("cam.jpg", np.uint8(255 * cam))

if __name__ == '__main__':

    trans = transforms.ToTensor()
    pp = np.loadtxt('model2_img1.txt')
    ct_img = cv2.imread('1.png', 1)

    img1 = cv2.resize(pp, (224, 224))
    img2 = np.float32(cv2.resize(ct_img, (224, 224))) / 255
    ct = img2
    img1 = gaussian_filter(img1).transpose(2,0,1)[0]

    means = [0.5]
    stds = [1]
    preprocessed_img = img1.copy()[:, :]
    preprocessed_img[:,:] = preprocessed_img[:, :] - means
    preprocessed_img[:,:] = preprocessed_img[:, :] / stds
    preprocessed_img = np.ascontiguousarray(preprocessed_img)
    preprocessed_img = torch.from_numpy(preprocessed_img)
    input1 = preprocessed_img
    heatmap = cv2.applyColorMap(np.uint8(255 * input1), cv2.COLORMAP_JET)
    heatmap = np.float32(heatmap) / 255

    heatmap = torch.where(torch.from_numpy(heatmap) < 0.6, torch.tensor([0.0]), torch.from_numpy(heatmap))
    cam = heatmap + np.float32(img2)
    cam = cam / cam.max()

    cv2.imwrite("cam.jpg", np.uint8(255 * cam))

    cam = torch.where(torch.from_numpy(heatmap) < 0.005, torch.from_numppy(img2), torch.from_numpy(heatmap+np.float32(img2)))


import librosa
import numpy as np
import pandas as pd
from IPython.display import Audio
import os
import shutil
import skimage
import skimage.io
from PIL import Image

class Preprocess:
    def __init__(self, hop_length=256, n_mels=80, win_size=128):
        self.hop_length = hop_length
        self.n_mels = n_mels
        self.n_fft = hop_length * 4
        self.frame_length = hop_length * 4
        self.win_size= win_size

    def scale_minmax(self, X, min=0, max=255):
        '''
        Scales and converts the spectrogram to a PNG
        Called by spectrogram_image()
        '''
        X_std = (X - X.min()) / (X.max() - X.min())
        X_scaled = X_std * (max - min) + min
        return X_scaled

    def loudest_window(self, spectrogram):
        '''
        Identifies the loudest window of a specified size within a spectrogram
        '''
        # Convert the array to a spectrogram
        array = np.asarray(spectrogram)
        # get the width of the array
        width = array.shape[1]
        # set the starting and ending of the window
        start = 0
        end = self.win_size
        arrays = []
        # move the window across the image advancing by 1 pixel.
        # append the array in the window to the list of arrays
        while end <= width:
            window = array[0:self.n_mels,start:end]
            arrays.append(window)
            start += 1
            end += 1
        # get the sum of each window in the array
        sums = [np.sum(window) for window in arrays]
        # get the index of the smallest window. This will be the loudest section
        index = sums.index(max(sums))
        # pull the loudest window from the list and return it.
        loudest = arrays[index]
        return loudest

    def add_silence(self, spectrogram):
        '''
        Adds silence to the end of a spectrogram if the spectrogram is shorter than
        the minimum window size
        '''
        # get the width of the spectrogram
        width = spectrogram.shape[1]
        # subtract the width from the minimum width to determine number of silent columns to add
        n_col = self.win_size-width
        # create silent columns
        silence = np.full((self.n_mels, n_col), 0)
        # convert the spectrogram to an array
        spect_ar = np.asarray(spectrogram)
        # add the silence to the spectrogram array
        spect_ar = np.concatenate((spect_ar, silence), axis = 1).astype(np.uint8)
        return spect_ar

    def standardize(self, spectrogram):
        '''
        Standardizes an image using the loudest_window() and add_silence() functions
        Pass spectrogram as an array
        '''
        width = spectrogram.shape[1]
        # if the width of the spectrogram is less than the minimum, add silence
        if width < self.win_size:
            spect = self.add_silence(spectrogram)
        # otherwise get the loudest window
        else:
            spect = self.loudest_window(spectrogram)
        # scale and return
        return scale_minmax(spect).astype(np.uint8)

    def spectrogram_image(self, wav_file, out_file = '', save = True):
        '''
        Converts a WAV file to a spectrogram and sames the image

        wav_file: The path to a WAV file as a string
        '''
        # load the audio data
        y, sr = librosa.load(wav_file)
        # trim the silence at edges
        #sample, _ = librosa.effects.trim(y, top_db=30, ref = np.max, hop_length = 128, frame_length=512)
        # use log-melspectrogram
        mels = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=self.n_mels,
                                                n_fft=self.hop_length*4, hop_length=self.hop_length)
        # convert mel to decibles
        DB = librosa.power_to_db(mels, ref=np.max)
        # Standardize by finding the loudest window or adding silence
        # and rescaling
        DB = self.standardize(DB)

        # save the image
        if save == True:
            skimage.io.imsave(out_file, DB)
        else:
            return Image.fromarray(DB, mode = 'P')

    def play_spectrogram(self, spectrogram, sr = 22050, n_fft = 1024, hop_length = 256):
        array = np.asarray(spectrogram)
        mels = librosa.db_to_power(array, ref=1)
        return Audio(librosa.feature.inverse.mel_to_audio(mels, sr = sr, n_fft = self.hop_length*4, hop_length = self.hop_length), rate = sr)


import torch
from PIL import Image
from torchvision import transforms
import torchvision
from tqdm import tqdm
import argparse
from efficientnet_pytorch import EfficientNet
from onnxsim import simplify
import onnx
import rexnetv1
import osnet
from thop import profile, clever_format
import numpy as np

def argparser():
    parser = argparse.ArgumentParser(description='Measure Accuracy in ImageNet Val')
    parser.add_argument('--model', default='resnet18', choices=['resnet18','resnet34','resnet50','resnet101', 'resnet152','efficientnet_b0','efficientnet_b1','efficientnet_b2','efficientnet_b3','efficientnet_b4','rexnetv1_1.0','rexnetv1_1.3','rexnetv1_1.5','rexnetv1_2.0','mobilenetv2','mobilenetv3-s','mobilenetv3-l', 'osnet_1.0','osnet_0.75','osnet_0.5', 'osnet_0.25'],
                        help='Detector model name')
    parser.add_argument('--export', default=False,
                        action='store_true',
                        help='dump network to onnx format')
    parser.add_argument('--data_path', default='/home/ysjo/dataset/dataset/imagenet/val/', type=str,
                        help='imagenet val dataset path')
    parser.add_argument('--batch_size', default=64, type=int,
                        help='Batch size for training')
    parser.add_argument('--num_workers', default=16, type=int,
                        help='Number of workers used in dataloading')
    args = parser.parse_args()

    return args

def accuracy(outputs, targets):
    """Compute top-5 top-1"""
    _, pred = outputs.topk(5, dim=1, largest=True, sorted=True)
    y_ = targets.unsqueeze(1).expand_as(pred)
    correct = (pred == y_).squeeze()
    correct = correct.detach().cpu()
    correct_1 = torch.sum(correct[:,:1]).item()
    correct_5 = torch.sum(correct[:,:5]).item()

    return correct_1, correct_5

if __name__ == "__main__":
    args = argparser()

    # select model
    # resnet
    if args.model == 'resnet18':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
        input_sizes = (224,224)
    elif args.model == 'resnet34':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet34', pretrained=True)
        input_sizes = (224,224)
    elif args.model == 'resnet50':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
        input_sizes = (224,224)
    elif args.model == 'resnet101':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet101', pretrained=True)
        input_sizes = (224,224)
    elif args.model == 'resnet152':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet152', pretrained=True)
        input_sizes = (224,224)
    # efficientnet
    elif args.model == 'efficientnet_b0':
        model = EfficientNet.from_pretrained('efficientnet-b0')
        model.set_swish(memory_efficient=False)
        image_size = EfficientNet.get_image_size('efficientnet-b0') # 224
        input_sizes = (image_size,image_size)
    elif args.model == 'efficientnet_b1':
        model = EfficientNet.from_pretrained('efficientnet-b1')
        model.set_swish(memory_efficient=False)
        image_size = EfficientNet.get_image_size('efficientnet-b1') # 224
        input_sizes = (image_size,image_size)
    elif args.model == 'efficientnet_b2':
        model = EfficientNet.from_pretrained('efficientnet-b2')
        model.set_swish(memory_efficient=False)
        image_size = EfficientNet.get_image_size('efficientnet-b2') # 224
        input_sizes = (image_size,image_size)
    elif args.model == 'efficientnet_b3':
        model = EfficientNet.from_pretrained('efficientnet-b3')
        model.set_swish(memory_efficient=False)
        image_size = EfficientNet.get_image_size('efficientnet-b3') # 224
        input_sizes = (image_size,image_size)
    elif args.model == 'efficientnet_b4':
        model = EfficientNet.from_pretrained('efficientnet-b4')
        model.set_swish(memory_efficient=False)
        image_size = EfficientNet.get_image_size('efficientnet-b4') # 224
        input_sizes = (image_size,image_size)
    # rexnet
    elif args.model == 'rexnetv1_1.0':
        model = rexnetv1.ReXNetV1(width_mult=1.0)
        model.load_state_dict(torch.load('./rexnetv1_1.0.pth'))
        input_sizes = (224,224)
    elif args.model == 'rexnetv1_1.3':
        model = rexnetv1.ReXNetV1(width_mult=1.3)
        model.load_state_dict(torch.load('./rexnetv1_1.3.pth'))
        input_sizes = (224,224)
    elif args.model == 'rexnetv1_1.5':
        model = rexnetv1.ReXNetV1(width_mult=1.5)
        model.load_state_dict(torch.load('./rexnetv1_1.5.pth'))
        input_sizes = (224,224)
    elif args.model == 'rexnetv1_2.0':
        model = rexnetv1.ReXNetV1(width_mult=2.0)
        model.load_state_dict(torch.load('./rexnetv1_2.0.pth'))
        input_sizes = (224,224)
    # mobilenet
    elif args.model == 'mobilenetv2':
        model = torch.hub.load('pytorch/vision:v0.10.0', 'mobilenet_v2', pretrained=True)
        input_sizes = (224,224)
    elif args.model == 'mobilenetv3-s':
        model = torchvision.models.mobilenet_v3_small(weights='DEFAULT')
        input_sizes = (224,224)
    elif args.model == 'mobilenetv3-l':
        model = torchvision.models.mobilenet_v3_large(weights='DEFAULT')
        input_sizes = (224,224)
    # osnet
    elif args.model == 'osnet_1.0':
        model = osnet.osnet_x1_0()
        input_sizes = (224, 224)
    elif args.model == 'osnet_0.75':
        model = osnet.osnet_x0_75()
        input_sizes = (224, 224)
    elif args.model == 'osnet_0.5':
        model = osnet.osnet_x0_5()
        input_sizes = (224, 224)
    elif args.model == 'osnet_0.25':
        model = osnet.osnet_x0_25()
        input_sizes = (224, 224)
    else:
        raise ValueError("no model")
    input = torch.randn(1, 3, input_sizes[0],  input_sizes[1])
    macs, params = profile(model, inputs=(input, ))
    macs, params = clever_format([macs, params], "%.3f")
    model.eval()
    print(f" INPUT SIZE : {input_sizes} ")

    # export model
    if args.export:
        dummy_input = torch.randn(1, 3, input_sizes[0], input_sizes[1])
        torch.onnx.export(model, dummy_input, args.model+".onnx", verbose=False, input_names=['images'], output_names=['output'])
        onnx.checker.check_model(args.model+".onnx")
        model_onnx, check = simplify(args.model+".onnx")
        onnx.save(model_onnx, args.model+".onnx")

    if args.model[:12] == 'efficientnet':
        preprocess = transforms.Compose([
            transforms.Resize(input_sizes[0]+32),
            transforms.CenterCrop(input_sizes),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    else:
        preprocess = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(input_sizes),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    testset=torchvision.datasets.ImageFolder(root=args.data_path, transform=preprocess)
    testloader=torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True, drop_last=False)

    # move the input and model to GPU for speed if available
    if torch.cuda.is_available():
        model = model.to('cuda')

    top_1 = 0
    top_5 = 0
    p_bar = tqdm(testloader)
    with torch.no_grad():
        for data in p_bar:
            imgs, labels = data
            if torch.cuda.is_available():
                imgs = imgs.to('cuda')
                labels = labels.to('cuda')
            outputs = model(imgs)
            probabilities = torch.exp(outputs)
            ret1, ret2 = accuracy(probabilities, labels)
            top_1 += ret1
            top_5 += ret2

    # for measure inference time average 100 repetitions
    repetitions = 100
    timings = np.zeros((repetitions,1))
    starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
    img = imgs[0].unsqueeze(0)   # 1 image
    with torch.no_grad():
        for rep in range(repetitions):
            starter.record()
            _ = model(imgs)
            ender.record()
            # WAIT FOR GPU SYNC
            torch.cuda.synchronize()
            curr_time = starter.elapsed_time(ender)
            timings[rep] = curr_time
    print("!! Sumary !!")
    print(f"MACs : {macs}")
    print(f"params : {params}")
    print(f"FLOPS : {float(macs[:-1]) *2 }{macs[-1]}")
    print("!! average Inference Time !!")
    print(f"Inference time : {np.mean(timings)}ms")
    print("!! Accuracy !!")
    print(f"top-1 : {round(top_1/500, 2)}")
    print(f"top-5 : {round(top_5/500, 2)}")
    print(f"top-1-error : {round(100 - top_1/500,2)}")
    print(f"top-5-error : {round(100 - top_5/500,2)}")
    print("!! DONE !!")


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import pybullet as p
import tf.transformations as tt
import matplotlib.pyplot as plt
from scipy.stats import norm
import time

def addDebugFrame(pos, quat, length=0.2, lw=2, duration=0.2, pcid=0):
    rot_mat = np.array(p.getMatrixFromQuaternion(quat)).reshape(3, 3).transpose()
    toX = pos + length * rot_mat[0, :]
    toY = pos + length * rot_mat[1, :]
    toZ = pos + length * rot_mat[2, :]
    item_id = [0, 0, 0]
    item_id[0] = p.addUserDebugLine(pos, toX, [1, 0, 0], lw, duration, physicsClientId=pcid)
    item_id[1] = p.addUserDebugLine(pos, toY, [0, 1, 0], lw, duration, physicsClientId=pcid)
    item_id[2] = p.addUserDebugLine(pos, toZ, [0, 0, 1], lw, duration, physicsClientId=pcid)
    return item_id

class CameraPose(object):

    def __init__(self, position, rpy, distance):
        self.position = position
        self.rpy = rpy
        self.distance = distance

def normal_random_integer(samples = 1, mean = 0, std = 1):

    return norm.ppf(np.random.random(samples), loc=mean, scale=std).astype(int)

class ImageGetter(object):
    def __init__(self, pcid):
        self._id = pcid

        self.camera_poses = []
        self.camera_poses.append(CameraPose([0.3, 0.0, 0.5],[0.0, -30.0, 90.0],1.5))
        self.camera_poses.append(CameraPose([0.50,0.35,0.20],[0.0, 4.60, 168.0],1.5))
        self.camera_poses.append(CameraPose([1.0,-0.02,-1.31],[0.0, -46.6, 214.0],1.9))
        self.camera_poses.append(CameraPose([0.6,0.62,-0.55],[0.0, -39.0, 320.0],1.5))

        self.camera_idx = 0

        camera_pose = self.camera_poses[self.camera_idx]

        self.camTargetPos = camera_pose.position#[0.3, 0.0, 0.5]
        self.roll = camera_pose.rpy[0]
        self.pitch = camera_pose.rpy[1]
        self.yaw = camera_pose.rpy[2]

        self.camDistance = camera_pose.distance

        self.pixelWidth = 640
        self.pixelHeight = 640
        self.aspect = self.pixelWidth / self.pixelHeight

        self.nearPlane = 0.01
        self.farPlane = 5.0
        self.fov = 57.0
        # Primesense carmine 1.09 field of view
        #horizontal FOV of 57.5 degrees and vertical FOV of 45 degrees.

        self.xyz_rgb = np.zeros((self.pixelWidth * self.pixelHeight, 6))

        self.projectionMatrix = p.computeProjectionMatrixFOV(self.fov, self.aspect,
                                                             self.nearPlane, self.farPlane,
                                                             physicsClientId=self._id)
        self.viewMatrix = None

        self.setCameraPose()

    def setCameraPose(self, idx = 0):
        assert idx < len(self.camera_poses)

        self.camera_idx = idx

        camera_pose = self.camera_poses[self.camera_idx]

        self.camTargetPos = camera_pose.position
        self.roll = camera_pose.rpy[0]
        self.pitch = camera_pose.rpy[1]
        self.yaw = camera_pose.rpy[2]
        self.camDistance = camera_pose.distance

        self.viewMatrix = p.computeViewMatrixFromYawPitchRoll(self.camTargetPos, self.camDistance,
                                                              yaw=self.yaw, pitch=self.pitch, roll=self.roll,
                                                              upAxisIndex=2, physicsClientId=self._id)

        p.resetDebugVisualizerCamera(cameraDistance=self.camDistance,
                                    cameraYaw=self.yaw, cameraPitch=self.pitch,
                                    cameraTargetPosition = self.camTargetPos
                                    )

        self.setCameraParams()

    def nextCameraPose(self):
        idx = (self.camera_idx+1)%len(self.camera_poses)

        self.setCameraPose(idx)

    def previousCameraPose(self):
        idx = (self.camera_idx-1)
        idx = (len(self.camera_poses)-1) if idx < 0 else idx

        self.setCameraPose(idx)

    def setCameraParams(self):
        # https://blog.noctua-software.com/opencv-opengl-projection-matrix.html
        opengl_proj_mat = np.array(self.projectionMatrix).reshape([4,4])
        opengl_view_mat = np.array(self.viewMatrix).reshape([4,4])

        self.fx = opengl_proj_mat[0,0] / 2.0 * self.pixelWidth
        self.fy = -opengl_proj_mat[1,1] / 2.0 * self.pixelHeight
        self.cx = (1.0 - opengl_proj_mat[2,0])/2.0*self.pixelWidth
        self.cy = (1.0 + opengl_proj_mat[2,1])/2.0*self.pixelHeight

        self.extrinsic_params = np.linalg.inv(opengl_view_mat).T
        self.camera_pos = self.extrinsic_params[0:3, 3]
        self.camera_quat = tt.quaternion_from_matrix(self.extrinsic_params)

        addDebugFrame(self.camera_pos, self.camera_quat, 0.2, 2, 10.0, self._id)

    def getImage(self):
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING,0)
        Renderer = p.ER_BULLET_HARDWARE_OPENGL #p.ER_TINY_RENDERER
        img_arr = p.getCameraImage(self.pixelWidth, self.pixelHeight,
                                   self.viewMatrix, self.projectionMatrix,
                                   renderer=Renderer, shadow=-1, physicsClientId=self._id)

        self.rgb = img_arr[2] # RGBA data
        self.dep = img_arr[3] # depth data
        p.configureDebugVisualizer(p.COV_ENABLE_RENDERING, 1)

    def corruptImage(self, depth_image):
        # zˆ(u, v) = z(u + N (0, σ2p), v + N (0, σ2p)) + N (0, σ2d)
        # Ed John's paper
        # Corrupt image with noise
        rows, cols = depth_image.shape

        std_dev = 0.001
        std_p = 1
        std_v = 1
        rintc = normal_random_integer(rows*cols,0,std_p).reshape(rows,cols)
        rintr = normal_random_integer(rows*cols,0,std_p).reshape(rows,cols)

        c, r = np.meshgrid(np.arange(cols), np.arange(rows), sparse=False)
        c_new = np.minimum(np.maximum(rintc + c,0),rows-1)
        r_new = np.minimum(np.maximum(rintr + r,0),cols-1)

        depth_image[c,r] = depth_image[c_new,r_new] # zˆ(u, v) = z(u + N (0, σ2p), v + N (0, σ2p))

        noise = np.random.randn(rows, cols)*std_dev
        noise_cutoff = 0.005
        noise[ noise < -noise_cutoff ] = -noise_cutoff
        noise[ noise > noise_cutoff ] = noise_cutoff
        depth_image = depth_image + noise # zˆ(u, v) = zˆ(u, v) + N (0, σ2d)

        # mask = (np.random.uniform(0,1,size=depth_image.shape) > 0.95).astype(np.bool)
        # # Randomly set 5% of the depth pixels to 0
        # depth_image[mask] = 0

        return depth_image

    def createPointCloud(self):
        xyz_rgb = self.xyz_rgb

        depth_image = np.array(self.dep).T
        depth_image = self.farPlane* self.nearPlane / (self.farPlane - (self.farPlane - self.nearPlane) * depth_image)
        # i = 0

        rows, cols = depth_image.shape

        t0 = time.time()

        depth_image = self.corruptImage(depth_image)

        # plt.imshow(depth_image)
        # plt.show()

        c, r = np.meshgrid(np.arange(cols), np.arange(rows), sparse=True)
        valid = (depth_image > 0) #
        z = np.where(valid, -depth_image , np.nan)
        x = np.where(valid, -z * (r - self.cx) / self.fx, 0)
        y = np.where(valid, -z * (c - self.cy) / self.fy, 0)
        # rgb = np.flipud(np.rot90(self.rgb))#[:, :,:3] #np.where(valid, self.rgb[r, c,:], [0,0,0])
        #
        # points = np.dstack((x, y, z,rgb))
        points = np.dstack((x, y, z))
        point_list = points.reshape(-1,3)
        # point_list = points.reshape(-1,7)
        # print "RGB shape: ", rgb.shape
        # points = np.vstack([x,y,z])
        # print "Point shape: ", points.shape
        # points = np.vstack([x,y,z,rgb])#np.dstack((x, y, z, rgb))
        # print points.shape

        # colors = np.where(valid, rgb[col, row, :3], 0)
        # points = np.dstack((x, y, z))
        # for row in range(self.pixelWidth):
        #     for col in range(self.pixelHeight):
        #         z = depth_image[row, col]
        #         xyz_rgb[i, 0] = (row - self.cx) / self.fx * z
        #         xyz_rgb[i, 1] = (col - self.cy) / self.fy * z
        #         xyz_rgb[i, 2] = -z
        #         xyz_rgb[i, 3:6] = self.rgb[col, row, :3]
        #         i = i + 1
        # print "Got it!"
        print("Time elapsed %.3f"%(time.time()-t0,))
        return point_list