from langchain.prompts import PromptTemplate
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import CTransformers
from langchain.chains import RetrievalQA #This is just a retrieval chain, for chat history use conversational retrieval chain
from langchain.chains import ConversationalRetrievalChain

from typing import Dict, Any
import torch

from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import ingest
import os

# LLMChain: This chain uses a Language Model for generating responses to queries or prompts.
# It can be used for various tasks such as chatbots, summarization, and more

# StuffDocumentsChain: This is the default QA chain used in the RetrievalQAChain. It processes the retrieved documents
# and generates answers to questions based on the content of the documents

#RetrievalQAChain: This chain combines a Retriever and a QA chain.
#It is used to retrieve documents from a Retriever and then use a QA chain to answer a question based on the retrieved documents

# the RetrievalQAChain is used with a VectorStore as the Retriever and the default StuffDocumentsChain as the QA chain.

#From what I understand:
# RetrievalQAChain: retrieves the vector based on the prompt
# StuffDocumentsChain: Converts the vector into the answer
# LLMChain: Uses the answer to generate a response to the prompt

DEVICE = "cpu"
if torch.cuda.is_available():
    DEVICE = "cuda"

DB_FAISS_PATH = "vectorstores/db_faiss"

custom_prompt_template = """Use the following pieces of information to answer the user's question.
If you don't know the answer, please just say that you don't know the answer, don't try to make up an answer.
Be empathetic, sympathetic, and kind in your responses.

Context: {chat_history}
{context}

Question: {question}

Only returns the helpful answer below and nothing else.
Helpful answer:
"""

custom_template = """Given the following conversation and a follow up question,
rephrase the follow up question to be a standalone question.
Preserve the original question in the answer sentiment during rephrasing.

Chat History:
{chat_history}
Follow Up Input: {question}
Standalone question:"""
CONDENSE_QUESTION_PROMPT_CUSTOM = PromptTemplate(template=custom_template, input_variables=["question", "chat_history"])

def setCustomPrompt():
    prompt = PromptTemplate(template=custom_prompt_template, input_variables=["context", "question", "chat_history"])
    return prompt

def loadLLM():
    llm = CTransformers(
        model="llama-2-7b-chat.ggmlv3.q8_0.bin", model_type="llama", max_new_tokens=512, temperature=0)
    return llm

def retrievalQAChain(llm, prompt, db):
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        chain_type="stuff",
        retriever= db.as_retriever(search_kwargs={"k": 1}),
        combine_docs_chain_kwargs={"prompt": prompt},
        return_source_documents = True,
        verbose=True,
        # condense_question_prompt=CONDENSE_QUESTION_PROMPT_CUSTOM,
        rephrase_question = False
        )
    #search_kwargs={"k": 2} means 2 searches
    #return_source_documents = True means don't use base knowledge use only knowledge we provided
    return qa_chain

def qaBot():
    embeddings = HuggingFaceEmbeddings(model_name = 'sentence-transformers/all-MiniLM-L6-v2', model_kwargs={"device": DEVICE})
    db = FAISS.load_local(DB_FAISS_PATH, embeddings)

    llm = loadLLM()
    qa_prompt = setCustomPrompt()
    qa = retrievalQAChain(llm, qa_prompt, db)

    return qa

def finalResult(query):
    qa_result = qaBot()
    chat_history = []
    # Will be query if using RetrievalQA, question for ConversationalQA
    response = qa_result({'chat_history': chat_history, 'question': query})
    print()
    return response

def askBot(query, qa_result, chat_history):
    # Will be query if using RetrievalQA, question for ConversationalQA
    response = qa_result({'chat_history': chat_history, 'question': query})
    print()
    return response

app = Flask(__name__)
app.static_folder = "static"
CORS(app)

qa_result = None
if os.path.exists(DB_FAISS_PATH):
    qa_result = qaBot()
chat_history = []

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

@app.route("/get")
def get_bot_response():
    global qa_result
    prompt = request.args.get("msg")
    answer = askBot(prompt, qa_result, chat_history)
    data = {
        "answer" : answer["answer"],
        "source" : answer["source_documents"][0].metadata["source"]
    }
    return jsonify(data)

@app.route("/get2")
def get_bot_response2():
    global qa_result
    link = request.args.get("msg")
    data = {}
    if ingest.addLink(link):
        ingest.createVectorDB(link)
        data = {
            "answer" : "Success!"
        }
    else:
        data = {
            "answer" : "Failed :("
        }
    qa_result = qaBot()
    return data

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=5000, debug=True)

    #Need to do Add button before send if the vector db path does not exist


from rest_framework import mixins
from api.gateway.serializers import GatewayGetSerializer, GatewayWriteSerializer

from common.views.common_viewset import CommonViewSet
from gateway.models import Gateway

class GatewayViewSet(
    mixins.CreateModelMixin,
    mixins.ListModelMixin,
    mixins.RetrieveModelMixin,
    mixins.DestroyModelMixin,
    CommonViewSet,
):
    queryset = Gateway.objects.all()

    serializer_class = {
        "create": GatewayWriteSerializer,
        "list": GatewayGetSerializer,
        "retrieve": GatewayGetSerializer,
    }


#!/usr/bin/env python

import digital_rf as drf
import matplotlib.pyplot as plt
import numpy as n
import stuffr

# real-time plot, just plot one integration period at last data in directory.
import sys
ch="ch1"
if len(sys.argv) > 1:
    print("using channel %s"%(sys.argv[1]))
    ch=sys.argv[1]

d=drf.DigitalRFReader("/data0/test200e3_7.953e6")

print("Recorded channels")
print(d.get_channels())
prop=d.get_properties(ch)
b=d.get_bounds(ch)

sr=prop["samples_per_second"]

i0=long(n.floor(b[1]/sr))*sr

ipp=100000/5
z=d.read_vector_c81d(i0,ipp,ch)
plt.plot(z.real)
plt.plot(z.imag)
plt.title("t=%s (UTC)"%(stuffr.unix2datestr(i0/sr)))
plt.show()


import sys, os
sys.path.insert(0, os.path.abspath('..'))

import numpy as np
import math
import scipy.linalg
from pypsmove.transformations import quaternion_matrix, quaternion_conjugate, quaternion_multiply
from filterpy.common import Q_discrete_white_noise, dot3
from filterpy.kalman import MerweScaledSigmaPoints
from filterpy.kalman.unscented_transform import unscented_transform
from filterpy.kalman.UKF import UnscentedKalmanFilter as UKF

# Load the test data
trainingdata = np.genfromtxt(os.path.join('..', '..', 'test_data', 'stationary.csv'), delimiter=',')  # Recorded while controller was stationary upright
testdata = np.genfromtxt(os.path.join('..', '..', 'test_data', 'movement.csv'), delimiter=',')  # Recorded while controller was moving around.
mean_dt = np.nanmean(np.diff(trainingdata[:, -1]))

# These were the calibration values taken from PSMoveService when stationary.csv and movement.csv were collected.
calibration = {
    'Accel': {
        'X': {'k': 0.000228754434, 'b': -0.0156696439},
        'Y': {'k': 0.000230096644, 'b': 0.0416475534},
        'Z': {'k': 0.00023118715, 'b': -0.0112125874},
        'NoiseRadius': 0.00866362546
    },
    'Gyro': {
        'X': {'k': 0.00152541522, 'b': 0.0},
        'Y': {'k': 0.00177830202, 'b': 0.0},
        'Z': {'k': 0.00164234079, 'b': 0.0},
        'Variance': 0.000382817292,
        'Drift': 0.0281366557
    },
    'Magnetometer': {
        'Center': {'X': -56, 'Y': 49, 'Z': -153},
        'BasisX': {'X': 1, 'Y': 0, 'Z': 0},
        'BasisY': {'X': 0, 'Y': 1, 'Z': 0},
        'BasisZ': {'X': 0, 'Y': 0, 'Z': 1},
        'Extents': {'X': 138, 'Y': 143, 'Z': 132},
        'Identity': {'X': 0.737549126, 'Y': 0.675293505, 'Z': 1},
        'Error': 54.9325409
    }
}

# DESIGN STATE VECTOR
#Define column indices of each state variable in state matrix.
st_pos_x, st_pos_y, st_pos_z = 0, 1, 2  # Position in 3D space in cm
st_lvel_x, st_lvel_y, st_lvel_z = 3, 4, 5  # Linear velocity in cm/s
st_lacc_x, st_lacc_y, st_lacc_z = 6, 7, 8  # Linear acceleration in cm/s^2
st_e_w, st_e_x, st_e_y, st_e_z = 9, 10, 11, 12  #Unit quaternion
st_avel_x, st_avel_y, st_avel_z = 13, 14, 15  # Angular velocity in rad/s
NSTATEDIM = 16
st_lin_list = (st_pos_x, st_pos_y, st_pos_z, st_lvel_x, st_lvel_y, st_lvel_z, st_lacc_x, st_lacc_y, st_lacc_z, st_avel_x, st_avel_y, st_avel_z)

# DESIGN STATE COVARIANCE VECTOR
p_pos_x, p_pos_y, p_pos_z = 0, 1, 2  # Position in 3D space in cm
p_lvel_x, p_lvel_y, p_lvel_z = 3, 4, 5  # Linear velocity in cm/s
p_lacc_x, p_lacc_y, p_lacc_z = 6, 7, 8  # Linear acceleration in cm/s^2
p_e_x, p_e_y, p_e_z = 9, 10, 11  #Angle-Axis representation
p_avel_x, p_avel_y, p_avel_z = 12, 13, 14  # Angular velocity in rad/s
NPDIM = 15
p_lin_list = (p_pos_x, p_pos_y, p_pos_z, p_lvel_x, p_lvel_y, p_lvel_z, p_lacc_x, p_lacc_y, p_lacc_z, p_avel_x, p_avel_y, p_avel_z)

# DESIGN OBSERVATION VECTOR
#Define the column indices of each observation variable in the observation matrix
obs_accel_x, obs_accel_y, obs_accel_z = 0, 1, 2  # In g-units
obs_gyro_x, obs_gyro_y, obs_gyro_z = 3, 4, 5  # In rad/s
obs_mag_x, obs_mag_y, obs_mag_z = 6, 7, 8  # In units of the strength of the local magnetic field.
obs_track_x, obs_track_y, obs_track_z = 9, 10, 11  # in cm
NOBSDIM = 12

# Some constants
GRAVITY = 9.806  # Depends on latitude.
grav_world = np.asarray([0, 0, -GRAVITY])  # -1 because accelerometer observations are returned normalized to gravity (I think)
LOCAL_MAG_FIELD = calibration['Magnetometer']['Identity']
mag_world = np.asarray([0, LOCAL_MAG_FIELD['X'], LOCAL_MAG_FIELD['Y'], LOCAL_MAG_FIELD['Z']])
QSigmaSq = 0.1
RScale = 5.0  # Only variable that appears to make any difference.
alpha= 0.1
beta= 2.0  # 2 is optimal for gaussian variables
kappa= 3 - NPDIM

# To transform vectors from controller frame to navigation frame...
#DCM = quaternion_matrix(q)
# To transform from navigation frame to controller frame...
#DCM.T = quaternion_matrix(q).T = quaternion_matrix(quaternion_conjugate(q))

def quat_from_axang(axang, dt=1):
    """
    Convert an axis-angle (3D) representation of orientation into a quaternion (4D)
    :param axang:
    :param dt:
    :return:
    """
    s = np.linalg.norm(axang)
    rot_angle = s * dt
    if rot_angle > 0:
        rot_axis = axang / s
        result = np.hstack((np.cos(rot_angle/2), rot_axis*np.sin(rot_angle/2)))
    else:
        result = np.asarray([1., 0., 0., 0.])
    return result / np.linalg.norm(result)

def angax_from_quat(quat):
    if quat[0] > 1:
        #if w > 1 acos and sqrt will produce errors, this cant happen if quaternion is normalised
        quat /= np.linalg.norm(quat)
    angle = 2 * math.acos(quat[0])
    s = math.sqrt(1 - quat[0]*quat[0])
    if (s < 0.001):
        axis = [1., 0., 0.]
    else:
        axis = quat[1:] / s
    return angle * (axis / np.linalg.norm(axis))

def shrink_state(state):
    # Handle the linear parts
    out = np.zeros(NPDIM)
    out[(p_lin_list),] = state[(st_lin_list),]
    # Handle the quaternion part
    quat = state[(st_e_w, st_e_x, st_e_y, st_e_z),]
    angax = angax_from_quat(quat)
    out[(p_e_x, p_e_y, p_e_z),] = angax
    return out

def pcol_to_state(pcol):
    out = np.zeros(NSTATEDIM)
    out[(st_lin_list),] = pcol[(p_lin_list),]
    out[(st_e_w, st_e_x, st_e_y, st_e_z),] = quat_from_axang(pcol[(p_e_x, p_e_y, p_e_z),])
    return out

def state_diff(x1, x2):
    """
    Calculate the difference between two state vectors (e.g., state and sigmapoint)
    Used in update calculation of state covariance (P): state_sp - state_mean
    Used in calculation of state-observation cross-covariance: [state_sp - state_mean][obs_sp - obs_mean]
    (Only used for the first part. Obs difference may use its own function)
    :param x1:
    :param x2:
    :return:
    """
    if len(x1) != len(x2):
        raise ValueError("Vectors must be the same length.")
    result = np.nan * np.ones(len(x1))

    # For linear measurements, this is pretty easy, just subtract
    result[st_lin_list,] = np.subtract(x1[st_lin_list,], x2[st_lin_list,])

    # For orientation it is more complicated.
    x2q_conj = quaternion_conjugate(x2[(st_e_w, st_e_x, st_e_y, st_e_z),])
    delta_q = quaternion_multiply(x1[(st_e_w, st_e_x, st_e_y, st_e_z),], x2q_conj)
    result[(st_e_w, st_e_x, st_e_y, st_e_z),] = delta_q / np.linalg.norm(delta_q)
    return result

def state_diff_p(x1, x2):
    return shrink_state(state_diff(x1, x2))

def state_covcol_diff(state, pcol):
    return state_diff(state, pcol_to_state(pcol))

def x_mean_fn(sigmas, Wm):
    # Calculate average state x from sigma points.
    x = np.nan*np.ones(sigmas.shape[1])
    # For position, velocity, acceleration, and angular velocity this is simple
    x[st_lin_list,] = np.dot(Wm, sigmas[:, st_lin_list])
    # For orientation, we need the weighted average quaternion.
    #http://stackoverflow.com/questions/12374087/average-of-multiple-quaternions
    Q = sigmas[:, (st_e_w, st_e_x, st_e_y, st_e_z)].T * Wm
    eig_val, eig_vec = np.linalg.eig(np.dot(Q, Q.T))
    q_out = eig_vec[:, np.argmax(eig_val)]
    q_out /= np.linalg.norm(q_out)
    x[(st_e_w, st_e_x, st_e_y, st_e_z),] = q_out
    return x

def process_fun(state, dt):
    """
    Describes how state at time k evolves to a state at time k+dt
    """
    output = np.nan * np.ones((NSTATEDIM,))

    # Position and linear updates
    output[(st_pos_x, st_pos_y, st_pos_z),] = state[(st_pos_x, st_pos_y, st_pos_z),]\
                                           + dt * state[(st_lvel_x, st_lvel_y, st_lvel_z),]\
                                           + dt * dt * state[(st_lacc_x, st_lacc_y, st_lacc_z),] / 2
    output[(st_lvel_x, st_lvel_y, st_lvel_z),] = state[(st_lvel_x, st_lvel_y, st_lvel_z),]\
                                              + dt*state[(st_lacc_x, st_lacc_y, st_lacc_z),]
    output[(st_lacc_x, st_lacc_y, st_lacc_z),] = state[(st_lacc_x, st_lacc_y, st_lacc_z),]

    # Orientation update
    delta_q = quat_from_axang(state[(st_avel_x, st_avel_y, st_avel_z),], dt)
    q_out = quaternion_multiply(state[(st_e_w, st_e_x, st_e_y, st_e_z),], delta_q)  #Maybe to do: *qnoise in the middle
    output[(st_e_w, st_e_x, st_e_y, st_e_z),] = q_out / np.linalg.norm(q_out)
    output[(st_avel_x, st_avel_y, st_avel_z),] = state[(st_avel_x, st_avel_y, st_avel_z),]

    return output

def observation_fun(state):
    output = np.nan * np.ones((NOBSDIM,))

    q = state[(st_e_w, st_e_x, st_e_y, st_e_z),]
    # world_to_cntrl = quaternion_matrix(quaternion_conjugate(q))
    accel_world_g = np.hstack((0, state[(st_lacc_x, st_lacc_y, st_lacc_z),] + grav_world)) / GRAVITY
    output[(obs_accel_x, obs_accel_y, obs_accel_z),] = quaternion_multiply(q, quaternion_multiply(accel_world_g, quaternion_conjugate(q)))[1:]
                                                        # + accel_bias + accel_noise
    output[(obs_gyro_x, obs_gyro_y, obs_gyro_z),] = state[(st_avel_x, st_avel_y, st_avel_z),]  # + gyro_bias + gyro_noise
    output[(obs_mag_x, obs_mag_y, obs_mag_z),] = quaternion_multiply(q, quaternion_multiply(mag_world,
                                                                                quaternion_conjugate(
                                                                                q)))[1:]
                                                # + mag_bias + mag_noise
    output[(obs_track_x, obs_track_y, obs_track_z),] = state[(st_pos_x, st_pos_y, st_pos_z),]
    return output

# Initial state
x_init = np.zeros((NSTATEDIM,))  # Assume all velocities and accelerations are initially 0
x_init[(st_pos_x, st_pos_y, st_pos_z),] = testdata[0, (obs_track_x, obs_track_y, obs_track_z)]  # First optical pos
x_init[(st_e_w, st_e_x, st_e_y, st_e_z),] = testdata[0, 12:16]  # First orientation. I could calc from accel/gyro, but it's already done for me...

# Initial state covariance
P_init = 0.1 * np.eye(NPDIM)  #TODO: calculate from trainingdata.

# Process noise covariance
Q_init = QSigmaSq*np.eye(NPDIM)
wn = Q_discrete_white_noise(3, mean_dt, var=QSigmaSq)  # 3rd order for position
Q_init[(p_pos_x, p_lvel_x, p_lacc_x),][:, (p_pos_x, p_lvel_x, p_lacc_x)] = wn
Q_init[(p_pos_y, p_lvel_y, p_lacc_y),][:,(p_pos_y, p_lvel_y, p_lacc_y)] = wn
Q_init[(p_pos_z, p_lvel_z, p_lacc_z),][:,(p_pos_z, p_lvel_z, p_lacc_z)] = wn
wn = Q_discrete_white_noise(2, mean_dt, var=QSigmaSq)  # Only 2nd order for orientation
Q_init[(p_e_x, p_avel_x),][:, (p_e_x, p_avel_x)] = wn
Q_init[(p_e_y, p_avel_y),][:, (p_e_y, p_avel_y)] = wn
Q_init[(p_e_z, p_avel_z),][:, (p_e_z, p_avel_z)] = wn

# Measurement noise covariance - calculate from trainingdata
# Every one of our measurements should be independent, so no off-diagonals.
R_init = RScale * np.diag(np.cov(trainingdata[:, :12], rowvar=False)) * np.eye(NOBSDIM)

#Subclass MerweScaledSigmaPoints and override sigma_points to get rid of size assertion
class MySigmaPoints(MerweScaledSigmaPoints):
    def sigma_points(self, x, P):
        n = self.n
        lambda_ = self.alpha**2 * (n + self.kappa) - n
        U = self.sqrt((lambda_ + n)*P)
        sigmas = np.zeros((2*n+1, x.size))
        sigmas[0] = x
        for k in range(n):
            sigmas[k+1]   = self.subtract(x, -U[k])
            sigmas[n+k+1] = self.subtract(x, U[k])
        return sigmas

class MyUKF(UKF):
    @property
    def P(self):
        return self._P
    @P.setter
    def P(self, val):
        self._P = val
        self._num_sigmas = 2 * self._P.shape[0] + 1
        self.sigmas_f = np.zeros((self._num_sigmas, self._dim_x))
        self.sigmas_h = np.zeros((self._num_sigmas, self._dim_z))

    def update(self, z, R=None, UT=None, hx_args=()):
        if z is None:
            return
        if not isinstance(hx_args, tuple):
            hx_args = (hx_args,)
        if UT is None:
            UT = unscented_transform
        if R is None:
            R = self.R
        elif np.isscalar(R):
            R = np.eye(self._dim_z) * R

        for i in range(self._num_sigmas):
            self.sigmas_h[i] = self.hx(self.sigmas_f[i], *hx_args)

        # mean and covariance of prediction passed through unscented transform
        zp, Pz = UT(self.sigmas_h, self.Wm, self.Wc, R, self.z_mean, self.residual_z)

        # compute cross variance of the state and the measurements
        Pxz = np.zeros((self.P.shape[0], self._dim_z))  # Changed
        for i in range(self._num_sigmas):
            dx = self.residual_x(self.sigmas_f[i], self.x)
            dz =  self.residual_z(self.sigmas_h[i], zp)
            Pxz += self.Wc[i] * np.outer(dx, dz)

        K = np.dot(Pxz, scipy.linalg.inv(Pz))   # Kalman gain
        y = self.residual_z(z, zp)   #residual
        self.x = state_covcol_diff(self.x, -np.dot(K, y))  # Changed
        self._P = self.P - dot3(K, Pz, K.T)

def myUT(sigmas, Wm, Wc, noise_cov=None,
                        mean_fn=None, residual_fn=None):
    kmax, n = sigmas.shape
    if mean_fn is None:
        # new mean is just the sum of the sigmas * weight
        x = np.dot(Wm, sigmas)    # dot = \Sigma^n_1 (W[k]*Xi[k])
    else:
        x = mean_fn(sigmas, Wm)
    if residual_fn is None:
        y = sigmas - x[np.newaxis,:]
        P = y.T.dot(np.diag(Wc)).dot(y)
    else:
        if noise_cov is not None:
            p_shape = noise_cov.shape
        else:
            y0 = residual_fn(sigmas[0], x)
            p_shape = (len(y0), len(y0))
        P = np.zeros(p_shape)
        for k in range(kmax):
            y = residual_fn(sigmas[k], x)
            P += Wc[k] * np.outer(y, y)
    if noise_cov is not None:
        P += noise_cov
    return (x, P)

# subtract wants non-sharnk version
sigpoints = MySigmaPoints(n=NPDIM, alpha=alpha, beta=beta, kappa=kappa, subtract=state_covcol_diff)

# residual_x wants shrank version
myukf = MyUKF(dim_x=NSTATEDIM, dim_z=NOBSDIM, dt=mean_dt,
            hx=observation_fun, fx=process_fun, points=sigpoints,
            x_mean_fn=x_mean_fn, residual_x=state_diff_p)
myukf.x = x_init
myukf.P = P_init
myukf.Q = Q_init
myukf.R = R_init

# Run the filter
observations = testdata[np.logical_not(np.any(np.isnan(testdata), axis=1))]
last_time = observations[0, -1] - mean_dt
uxs = []
for z in observations:
    dt = z[-1] - last_time
    myukf.predict(dt=dt, UT=myUT)
    myukf.update(z[:12], UT=None)
    uxs.append(myukf.x.copy())
    last_time = z[-1]
uxs = np.array(uxs)

# Plot the result
track_pos = observations[:, (obs_track_x, obs_track_y, obs_track_z)]
filt_pos = uxs[:, (st_pos_x, st_pos_y, st_pos_z)]
diffs = np.linalg.norm(filt_pos-track_pos, axis=1)
print("Difference mean: {}, std: {}".format(np.mean(diffs), np.std(diffs)))

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax1 = fig.add_subplot(211, projection='3d')
pl_tracker, = ax1.plot(track_pos[:, 0], track_pos[:, 1], track_pos[:, 2], 'b', label="Original")
pl_filter, = ax1.plot(filt_pos[:, 0], filt_pos[:, 1], filt_pos[:, 2], 'r', label="UKF")
ax1.legend(handles=[pl_tracker, pl_filter])

#TODO: Plot 3D end-point of quaternion on unit sphere. Show roll with different colour.
ax2 = fig.add_subplot(212)
psm_q, = ax2.plot(observations[:, -1], observations[:, 12], 'b', label="PSMoveService q0")
ukf_q, = ax2.plot(observations[:, -1], uxs[:, st_e_w], 'r', label="Predicted q0")
ax2.legend(handles=[psm_q, ukf_q])
plt.show()

#Additional TODOs:
#-handle partial observations (only IMU, only optical tracker)
#-compensate for sensor offset, lever arm
#-account for delays in optical measurements


# -*- coding: utf-8 -*-

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from sqlalchemy.ext.asyncio import AsyncSession
from backend.core import check_jwt_token
from backend.apis.deps import get_db
from utils.socket_msg import ConnectionManager
from backend.crud.message import getMessage

router = APIRouter()

manager = ConnectionManager()

@router.websocket(path="/message.io/{token:path}")
async def sys_message_io(
        websocket: WebSocket,
        token: str = Depends(check_jwt_token),
        db: AsyncSession = Depends(get_db),
):
    result = await getMessage.get_first(db)
    data = {"event": "get_unread_message", "message": "你有一条新的消息", "data": result}
    await manager.connect(websocket)
    await manager.broadcast(data)
    try:
        while True:
            get_msg = await websocket.receive_text()
            await manager.send_personal_message(data, websocket)
            await manager.broadcast(data)
    except WebSocketDisconnect as e:
        manager.disconnect(websocket)
        await manager.broadcast('账户登出')


#!/usr/bin/env python
# -*- coding=utf-8 -*-
# entity.py
# This file includes Thread and Session
# autor: Hang Liu, Dongao Li, Paerhati Remutula
# 10/2/17
#
import sys
import os
import socket
import struct
import random
import math
import codecs
from PyQt5.QtCore import (QCoreApplication, QObject, QRunnable, QThread,
                          QThreadPool, pyqtSignal)

from Crypto import Random
from Crypto.Hash import SHA
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.Cipher import AES
import base64
import uuid

from protocol import Protocol_Type

class EntityListenThread(QThread):
    hostInfo = ('', 0)
    secret_mode = False
    session_key = None

    # Signals
    setRemoteHostSignal = pyqtSignal('QString')
    receiveDataSignal = pyqtSignal(bytes, int)
    p_phase1_signal = pyqtSignal(bytes, bytes)
    p_phase2_signal = pyqtSignal(bytes, bytes, bytes)
    p_phase3_signal = pyqtSignal(bytes, bytes, bytes)
    p_phase4_signal = pyqtSignal(bytes)
    disconnect_signal = pyqtSignal()

    # decrypt file
    # This is converse function of encrypt file. It is similar to encrypt file in dialog.py
    # Using a tmp file to store received encrypted data.
    # shasum is received shasum from the other entity.
    def decrypt_file(self, filepath, shasum):
        f = open(filepath, 'wb')
        tmp = open('tmp', 'rb')
        h = SHA.new()
        while True:
            data = tmp.read(EntitySession.file_block_size)
            if not data:
                break
            decrypted = EntitySession.aesDecrypt(data, self.session_key)
            h.update(decrypted)
            f.write(decrypted)
        tmp.close()
        f.close()
        os.remove('tmp')
        print('received shasum: ', h.hexdigest())
        print('shasum: ', shasum)
        # Check received shasum and decrypted shasum. Are they equal or not
        if str.encode(h.hexdigest(), encoding='utf8') == shasum:
            # Shasums are equal
            return True
        else:
            # Shasums are not equal
            return False

    # This is thread entry. When user click the listening button, the function will be called
    def run(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(self.hostInfo)
            print(self.hostInfo)
            s.listen(5)
            protocol_type = Protocol_Type.TEXT
            while True:
                conn, addr = s.accept()
                self.setRemoteHostSignal.emit(addr[0])
                content_header_size = struct.calcsize('!I')
                buf = conn.recv(content_header_size)
                if buf:
                    filename = ''
                    filesize = 0
                    header = buf[0:content_header_size]
                    # Received message type, the value is number
                    temp, = struct.unpack('!I', header)
                    # Map the number of type to protocol enum
                    protocol_type = Protocol_Type(int(temp))
                    if protocol_type == Protocol_Type.FILE:
                        fileinfo_size = struct.calcsize('128sl')
                        content = conn.recv(fileinfo_size)
                        filename, filesize = struct.unpack('128sl', content)
                        fn = filename.strip(b'\00')
                        new_filename = os.path.join('./', 'new_' + bytes.decode(fn, encoding = "utf8"))
                        print ('file new name is %s, file size is %s'%(format(new_filename), filesize))
                        recvd_size = 0
                        fp = open('tmp', 'wb')
                        print ('start receiving...')

                        while not recvd_size == filesize:
                            data = ''
                            # Receive the data with 2048size
                            if filesize - recvd_size > EntitySession.file_block_size:
                                data = conn.recv(EntitySession.file_block_size)
                                recvd_size += len(data)
                            else:
                                # When the length of left data less than 2048, just receive left size data
                                data = conn.recv(filesize - recvd_size)
                                recvd_size = filesize

                            fp.write(data)
                        fp.close()
                        print ('end receive...')
                        if self.secret_mode:
                            data = conn.recv(40)
                            # using secure mode, need decrypt file first.
                            if self.decrypt_file(new_filename, data):
                                self.receiveDataSignal.emit(str.encode(new_filename, encoding='utf8'), 1)
                            else:
                                print('Some problems has been occured.')
                                pass
                        else:
                            # Rename tmp file to img file
                            os.rename('tmp', new_filename)
                            self.receiveDataSignal.emit(str.encode(new_filename, encoding='utf8'), 1)
                    # The process of message received
                    elif protocol_type == Protocol_Type.TEXT:
                        shead_size = struct.calcsize('l')
                        shead = conn.recv(shead_size)
                        content_size, = struct.unpack('l', shead)
                        if self.secret_mode:
                            content_size = content_size+40
                        content = conn.recv(content_size)
                        self.receiveDataSignal.emit(content, 0)
                    # The process of session creating
                    elif protocol_type == Protocol_Type.PROTOCOL:
                        # get header of wraper 1
                        p_header_size = struct.calcsize('!I64sl')
                        p_header = conn.recv(p_header_size)
                        # phase number, phase pattern, phase pattern size
                        p_phase, p_pattern, p_pattern_size = struct.unpack('!I64sl', p_header)

                        p_content = conn.recv(p_pattern_size)
                        p_pattern = p_pattern.strip(b'\00')

                        if int(p_phase) == 1:
                            # get header  of wraper 2
                            ds, pk = struct.unpack(p_pattern, p_content)
                            self.p_phase1_signal.emit(ds, pk)
                        elif int(p_phase) == 2:
                            # get header  of wraper 2
                            pk, confid, ds = struct.unpack(p_pattern, p_content)
                            self.p_phase2_signal.emit(pk, confid, ds)
                        elif int(p_phase) == 3:
                            # get header  of wraper 2
                            p_confid, ds, s_confid = struct.unpack(p_pattern, p_content)
                            self.p_phase3_signal.emit(p_confid, ds, s_confid)
                        elif int(p_phase) == 4:
                            # get header  of wraper 2
                            s_confid, = struct.unpack(p_pattern, p_content)
                            self.p_phase4_signal.emit(s_confid)
                    elif protocol_type == Protocol_Type.COMMAND:
                        shead_size = struct.calcsize('l')
                        shead = conn.recv(shead_size)
                        content_size, = struct.unpack('l', shead)
                        if self.secret_mode:
                            content_size = content_size+40
                        content = conn.recv(content_size)
                        if self.secret_mode:
                            content = EntitySession.aesDecrypt(content, self.session_key)
                            origin_content, flag = EntitySession.integrity(content)
                            if flag: # pass the verify of Integrity
                                content = origin_content
                            else: # not pass the verify of Integrity

                                pass
                        if content == b'disconnect':
                            self.disconnect_signal.emit()

                    conn.close()
            s.close()

class EntitySession(QObject):
    is_client = False
    is_server = False
    is_session_created = False
    using_secret = False
    session_id = ''
    nonce = 0
    remote_nonce = 0
    session_key = None
    random_generator = 0
    iv = b'1111111111111111'
    entity_uuid = ''
    private_key_name = 'entity-private.pem'
    public_key_name = 'entity-public.pem'
    glost_public_key_name = 'glost-public.pem'
    rsa_decrypt_err = 'Decrypt Error'
    local = ''
    remote = ''
    sendPhaseSignal = pyqtSignal(bytes)
    # record a size of data to process
    file_block_size = 2048
    def __init__(self, hostaddr, parent=None):
        QObject.__init__(self, parent)
        # get a id of A that will be stored database
        self.local = hostaddr
        self.entity_uuid = uuid.uuid5(uuid.NAMESPACE_DNS, hostaddr)

    # generate public and private key using rsa
    def generateRSAKeys(self):
        # generate psudo random number
        self.random_generator = Random.new().read
        # generate rsa
        rsa = RSA.generate(1024)

        # private key generating and writing as a file
        private_pem = rsa.exportKey()
        with open(self.private_key_name, 'w') as f:
            f.write(bytes.decode(private_pem, encoding='utf8'))
            f.close()
        # public key generating and writing as a file
        public_pem = rsa.publickey().exportKey()
        with open(self.public_key_name, 'w') as f:
            f.write(bytes.decode(public_pem, encoding='utf8'))
            f.close()

    # encrypt plaintext by public or private key
    def rsaEncrypt(self, plaintext, keyPath):
        key = ''
        with open(keyPath) as f:
            key = f.read()
            f.close()
        rsakey = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsakey)
        cipher_text = base64.b64encode(cipher.encrypt(str.encode(plaintext, encoding='utf8')))
        return cipher_text # This is a byte object

    # descrypt ciphertext by public or private key
    def rsaDecrypt(self, ciphertext, keyPath):
        key = ''
        with open(keyPath) as f:
            key = f.read()
            f.close()
        rsakey = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsakey)
        text = cipher.decrypt(base64.b64decode(ciphertext), self.rsa_decrypt_err)
        return text # This is a byte object

    # Important: pycrypto can not allowed using private key to encrypt message
    # explain as blow link
    # http://stackoverflow.com/questions/20057764/python-pycrypto-rsa-encrypt-method-gives-same-results-using-private-or-public-ke
    def rsaSign(self, plaintext, key, flag=False): #if flag=False, read key from file, else use it directly
        k = None
        if not flag:
            k = RSA.importKey(open(key).read())
        else:
            k = RSA.importKey(key)
        h = SHA.new(plaintext)
        signer = PKCS1_v1_5.new(k)
        return signer.sign(h)

    # Verify the signature of RSA
    def rsaVerifySign(self, plaintext, signature, key, flag=False):#if flag=False, read key from file, else use it directly
        k = None
        if not flag:
            k = RSA.importKey(open(key).read())
        else:
            k = RSA.importKey(key)
        h = SHA.new(plaintext)
        verifier = PKCS1_v1_5.new(k)
        if verifier.verify(h, signature):
            # print ("The signature is authenticated.")
            return True
        else:
            # print ("The signature is not authenticated.")
            return False

    # AES encrypt text
    @staticmethod
    def aesEncrypt(plaintext, key):
        cipher = AES.new(key, AES.MODE_CFB, EntitySession.iv)
        result = cipher.encrypt(plaintext)
        return result

    # AES decrypt text
    @staticmethod
    def aesDecrypt(ciphertext, key):
        cipher = AES.new(key, AES.MODE_CFB, EntitySession.iv)
        result = cipher.decrypt(ciphertext)
        return result

    # using sha to hash text
    @staticmethod
    def sha(text):
        h = SHA.new()
        h.update(str.encode(text, encoding='utf8'))
        return h.hexdigest()

    @staticmethod
    def sha_byte(text):
        h = SHA.new()
        h.update(text)
        return h.hexdigest()

    @staticmethod
    def hashText(text):
        hash_text = EntitySession.sha(text)
        print('Sent hash text:',hash_text)
        return hash_text + text

    @staticmethod
    def hashTextByte(text):
        hash_text = EntitySession.sha_byte(text)
        print('Sent hash text:',hash_text)
        return str.encode(hash_text, encoding='utf8') + text

    # Verify the integrity of message and command process
    @staticmethod
    def integrity(hash_text):
        hash_size = 40 # record the digitals of sha result
        hash_text_len = len(hash_text)
        text_len = hash_text_len - hash_size
        origin = hash_text[hash_size:]
        hashed = hash_text[0:hash_size]
        print('Received hash text:', hashed)
        if str.encode(EntitySession.sha_byte(origin), encoding='utf8') == hashed:
            return (origin, True)
        else:
            return (origin, False)

    def generateNonce(self, b):
        result = random.randint(1*pow(10, b-1), 10*pow(10, b-1)-1)
        return result

    def createSession(self):
        # generate asymmetric keys using RSA
        self.generateRSAKeys()
        if self.is_client and not self.is_server:
            self.protocol_phase1_generator()

    # A->B: {A}Ka-1, Ka+
    def protocol_phase1_generator(self):
        # A: uuid by address of A
        entity_identification = str.encode(str(self.entity_uuid), encoding='utf8')
        # {A}Ka-1
        digital_signature = self.rsaSign(entity_identification, self.private_key_name)
        # get Ka+ string
        public_key = ''
        with open(self.public_key_name) as f:
            public_key = f.read()
            f.close()
        # 1 means this is the phase 1
        pattern = '%ss%ss' % (len(digital_signature), len(str.encode(public_key, encoding='utf8')))
        p1_header = struct.pack('!I64sl', 1, str.encode(pattern, encoding='utf8'), struct.calcsize(pattern))
        msg = struct.pack(pattern, digital_signature, str.encode(public_key, encoding='utf8'))

        self.sendPhaseSignal.emit(p1_header+msg)
        print('===================1.The client has sent=================== \n Kc+, {H(C)}Kc-1')

    # B->A: Kb+, {sid, Nb}Ka+, {H(Nb)}Kb-1
    def protocol_phase2_generator(self):
        # Kb+
        public_key = ''
        with open(self.public_key_name) as f:
            public_key = f.read()
            f.close()
        # Nb
        self.nonce = self.generateNonce(10)
        # {sid, Nb}Ka+
        confiden_str = '||'.join([self.session_id, str(self.nonce)])
        confiden_result = self.rsaEncrypt(confiden_str, self.glost_public_key_name)

        # H(Nb) Dont need to hash nonce. When give a sign to nonce, sign function has hashed
        hash_nonce = EntitySession.sha(str(self.nonce))
        # {H(Nb)}Kb-1
        ds = self.rsaSign(str.encode(str(self.nonce), encoding='utf8'), self.private_key_name)

        # 2 means this is the phase 2
        pattern = '%ss%ss%ss' % (len(str.encode(public_key, encoding='utf8')), len(confiden_result), len(ds))
        p2_header = struct.pack('!I64sl', 2, str.encode(pattern, encoding='utf8'), struct.calcsize(pattern))
        msg = struct.pack(pattern, str.encode(public_key, encoding='utf8'), confiden_result, ds)
        self.sendPhaseSignal.emit(p2_header+msg)
        print('===================2.The server has sent=================== \n Ks+, {sid, Ns}Kc+, {H(Ns)}Kc-1 \n Ns = %s' % str(self.nonce))

    # A->B {Na}Kb+, {H(Na)}Ka-1, {Nb}Kab
    def protocol_phase3_generator(self):
        # Na
        self.nonce = self.generateNonce(10)
        # {Na}Kb+
        p_confid = self.rsaEncrypt(str(self.nonce), self.glost_public_key_name)
        # {H(Na)}Ka-1
        ds = self.rsaSign(str(self.nonce).encode('utf8'), self.private_key_name)

        # Generate Kab
        session_key = self.sha(str(self.nonce)+str(self.remote_nonce))
        # session key can has more 32 digits
        self.session_key = session_key[0:32]

        # {Nb}Kab
        s_confid = EntitySession.aesEncrypt(str.encode(str(self.remote_nonce), encoding='utf8'),
            str.encode(self.session_key, encoding='utf8'))

        # 3 means this is the phase 3
        pattern = '%ss%ss%ss' % (len(p_confid), len(ds), len(s_confid))
        p3_header = struct.pack('!I64sl', 3, str.encode(pattern, encoding='utf8'), struct.calcsize(pattern))
        msg = struct.pack(pattern, p_confid, ds, s_confid)
        self.sendPhaseSignal.emit(p3_header+msg)
        print('===================3.The client has sent=================== \n {Nc}Ks+, {H(Nc)}Kc-1, {Ns}Kcs \n Nc = %s' % str(self.nonce))

    # B->A: {Na}Kab
    def protocol_phase4_generator(self):
        s_confid = EntitySession.aesEncrypt(str(self.remote_nonce).encode('utf8'), self.session_key)
        # 4 means this is the phase 4
        pattern = '%ss' % (len(s_confid))
        p4_header = struct.pack('!I64sl', 4, str.encode(pattern, encoding='utf8'), struct.calcsize(pattern))
        msg = struct.pack(pattern, s_confid)
        self.sendPhaseSignal.emit(p4_header+msg)
        print('===================4.The server has sent=================== \n {Nc}Kcs}')

    # executed in B
    def protocol_phase1_resolver(self, hostaddr, ds, pk):
        glost_uuid = uuid.uuid5(uuid.NAMESPACE_DNS, hostaddr)
        self.saveGlostPublicKey(pk)
        if self.rsaVerifySign(str.encode(str(glost_uuid), encoding='utf8'), ds, pk, True):
            glost_uuid = uuid.uuid5(uuid.NAMESPACE_DNS, hostaddr)
        else:
            # re-connect
            pass

        # glost_uuid = self.rsaDecrypt(ds, self.glost_public_key_name)
        self.generateSessionID(glost_uuid)
        print('===================1.The server has received=================== \n Kc+, {H(C)}Kc-1')
        self.protocol_phase2_generator()

        # print(socket.gethostbyaddr(hostaddr))
    # executed in A
    def protocol_phase2_resolver(self, pk, confid, ds):
        info = bytes.decode(self.rsaDecrypt(confid, self.private_key_name), encoding='utf8')
        if info == self.rsa_decrypt_err:
            # do something or exit(1)
            pass
        else:
            self.session_id = info.split('||')[0]
            self.remote_nonce = int(info.split('||')[1])
            challenage = str(self.remote_nonce).encode('utf8')
            if self.rsaVerifySign(challenage, ds, pk, True):
                self.saveGlostPublicKey(pk)
            else:
                # re-send or re-resovle
                pass
            print('===================2.The client has received=================== \n Ks+, {sid, Ns}Kc+, {H(Ns)}Kc-1 \n Ns = %s' % self.remote_nonce)
            self.protocol_phase3_generator()

    # executed in B
    def protocol_phase3_resolver(self, p_confid, ds, s_confid):
        # Na
        remote_nonce = bytes.decode(self.rsaDecrypt(p_confid, self.private_key_name), encoding='utf8')
        if self.rsaVerifySign(remote_nonce.encode('utf8'), ds, self.glost_public_key_name):
            self.remote_nonce = int(remote_nonce)
        else:
            # re-send
            pass

        # Generate Kab
        session_key = self.sha(str(self.remote_nonce)+str(self.nonce))
        # session key can has more 32 digits
        i_session_key = session_key[0:32]

        # Nb
        nonce_return = bytes.decode(EntitySession.aesDecrypt(s_confid, i_session_key), encoding='utf8')
        if self.nonce == int(nonce_return):
            self.session_key = i_session_key
            is_session_created = True
            print('===================3.The server has received=================== \n {Nc}Ks+, {H(Nc)}Kc-1, {Ns}Kcs \n Nc = %s' % remote_nonce)
            self.protocol_phase4_generator()
        else:
            # re-send
            pass

    # executed in A
    def protocol_phase4_resolver(self, s_confid):
        nonce_return = bytes.decode(EntitySession.aesDecrypt(s_confid, self.session_key), encoding='utf8')
        if int(nonce_return) == self.nonce:
            is_session_created = True
        print('===================4.The client has received=================== \n {Nc}Kcs}')

    # save remote entity public key to a file
    def saveGlostPublicKey(self, key):
        with open(self.glost_public_key_name, 'w') as f:
            f.write(bytes.decode(key, encoding='utf8'))
            f.close()

    def generateSessionID(self, entity_uuid):
        self.session_id = str(entity_uuid)


from NeuralNetwork import NeuralNetwork
from Utils import load_input
import numpy as np
import matplotlib.pyplot as plt

dataset = load_input('../datasets/data.h5')
X_train, Y_train = dataset["X_train"]/255*2 - 1, dataset["Y_train"]
X_test, Y_test = dataset["X_test"]/255*2 - 1, dataset["Y_test"]

model = NeuralNetwork()
model.add_layer(20, 'relu', inputs=X_train.shape[0])
model.add_layer(7, 'relu')
model.add_layer(5, 'relu')
model.add_layer(1, "sigmoid")
model.compile(learning_rate=0.05, regularization='L2', lambd=0.01)

costs = model.fit(X_train, Y_train, 1000, verbose=True, step=10)

plt.plot(costs)
plt.show()

prediction_train = model.predict(X_train)
prediction_test = model.predict(X_test)

print("train accuracy: {} %".format(100 - np.mean(np.abs(prediction_train - Y_train)) * 100))
print("test accuracy: {} %".format(100 - np.mean(np.abs(prediction_test - Y_test)) * 100))


import matplotlib.pyplot as plt
import math
import particle as prt

dt_lista = []
aps_greska = []

p1 = prt.Particle()

dt = 0.000
for i in range(100):
    dt += 0.001
    p1.init(10,60,0,0,dt)
    greska = abs(p1.range() - p1.analiticki_domet())/p1.analiticki_domet()*100
    #print("Numericki: {:.2f}, analiticki: {:.2f}".format(p1.range(),p1.analiticki_domet()))
    #print("dt: {:.2f}, greska: {:.2f}".format(dt,greska))
    dt_lista.append(dt)
    aps_greska.append(greska)
    p1.reset()

plt.figure("Ovisnost relativne apsolutne pogreske o izboru vremenskog intervala")
plt.plot(dt_lista,aps_greska)
plt.xlabel("dt[s]")
plt.ylabel("apsolutna pogreska [%]")
plt.show()


from oslo_log import log as logging
from oslo_utils import excutils

from ironic.common import exception
from ironic.common.i18n import _
from ironic.common import utils as common_utils
from ironic.conductor import task_manager
from ironic.conductor import utils as manager_utils
from ironic.conf import CONF
from ironic.drivers import utils as driver_utils

LOG = logging.getLogger(__name__)

warn_about_default_boot_mode = False

def _set_boot_mode_on_bm(task, ironic_boot_mode, fail_if_unsupported=False):
    try:
        manager_utils.node_set_boot_mode(task, ironic_boot_mode)

    except exception.UnsupportedDriverExtension as ex:
        if fail_if_unsupported:
            msg = (_("Baremetal node %(uuid)s boot mode is not set "
                     "to boot mode %(boot_mode)s: %(error)s") %
                   {'uuid': task.node.uuid,
                    'boot_mode': ironic_boot_mode,
                    'error': ex})
            LOG.error(msg)
            raise exception.UnsupportedDriverExtension(msg)

        msg_tmpl = _("Baremetal node %(uuid)s boot mode is not set "
                     "to boot mode %(boot_mode)s. Assuming "
                     "baremetal node is already in %(boot_mode)s or "
                     "driver set boot mode via some other "
                     "mechanism: %(error)s")

        LOG.debug(msg_tmpl, {'uuid': task.node.uuid,
                             'boot_mode': ironic_boot_mode,
                             'error': ex})

    except exception.InvalidParameterValue as ex:
        msg = (_("Node %(uuid)s boot mode is not set. "
                 "Attempt to set %(ironic_boot_mode)s boot mode "
                 "on the baremetal node failed with error %(error)s") %
               {'uuid': task.node.uuid,
                'ironic_boot_mode': ironic_boot_mode,
                'error': ex})
        LOG.error(msg)
        raise exception.InvalidParameterValue(msg)

    else:
        LOG.info("Baremetal node boot mode is set to boot "
                 "mode %(boot_mode)s",
                 {'uuid': task.node.uuid, 'boot_mode': ironic_boot_mode})
        manager_utils.node_cache_boot_mode(task)

def sync_boot_mode(task):
    """Set node's boot mode from bare metal configuration

    Attempt to read currently set boot mode off the bare metal machine.
    Also read node's boot mode configuration:

    * If BM driver does not implement getting boot mode, assume
      BM boot mode is not set and apply the logic that follows
    * If Ironic node boot mode is not set and BM node boot mode is
      not set - set Ironic boot mode to `[deploy]/default_boot_mode`
    * If Ironic node boot mode is not set and BM node boot mode
      is set - set BM node boot mode on the Ironic node
    * If Ironic node boot mode is set and BM node boot mode is
      not set - set Ironic boot mode to BM boot mode
    * If both Ironic and BM node boot modes are set but they
      differ - try to set Ironic boot mode to BM boot mode and fail hard
      if underlying hardware type does not support setting boot mode

    In the end, the new boot mode may be set in
    'driver_internal_info/deploy_boot_mode'.

    :param task: a task object
    """
    node = task.node

    try:
        bm_boot_mode = manager_utils.node_get_boot_mode(task)

    except exception.UnsupportedDriverExtension as ex:
        bm_boot_mode = None

        LOG.debug("Cannot determine node %(uuid)s boot mode: %(error)s",
                  {'uuid': node.uuid, 'error': ex})

    ironic_boot_mode = get_boot_mode_for_deploy(node)

    # NOTE(etingof): the outcome of the branching that follows is that
    # the new boot mode may be set in 'driver_internal_info/deploy_boot_mode'

    if not ironic_boot_mode and not bm_boot_mode:
        default_boot_mode = CONF.deploy.default_boot_mode
        node.set_driver_internal_info('deploy_boot_mode', default_boot_mode)
        node.save()

        LOG.debug("Ironic node %(uuid)s boot mode will be set to default "
                  "boot mode %(boot_mode)s",
                  {'uuid': node.uuid, 'boot_mode': default_boot_mode})

        _set_boot_mode_on_bm(task, default_boot_mode)

    elif not ironic_boot_mode and bm_boot_mode:
        node.set_driver_internal_info('deploy_boot_mode', bm_boot_mode)
        node.save()

        LOG.debug("Ironic node %(uuid)s boot mode is set to boot mode "
                  "%(boot_mode)s reported by the driver",
                  {'uuid': node.uuid, 'boot_mode': bm_boot_mode})

    elif ironic_boot_mode and not bm_boot_mode:
        # NOTE(etingof): if only ironic boot mode is known, try to synchronize
        # (e.g. ironic -> bm) and do not fail if setting boot mode is not
        # supported by the underlying hardware type
        _set_boot_mode_on_bm(task, ironic_boot_mode)

    elif ironic_boot_mode != bm_boot_mode:
        msg = (_("Boot mode %(node_boot_mode)s currently configured "
                 "on node %(uuid)s does not match the boot mode "
                 "%(ironic_boot_mode)s requested for provisioning. "
                 "Attempting to set node boot mode to %(ironic_boot_mode)s.") %
               {'uuid': node.uuid, 'node_boot_mode': bm_boot_mode,
                'ironic_boot_mode': ironic_boot_mode})
        LOG.info(msg)

        # NOTE(etingof): if boot modes are known and different, try
        # to synchronize them (e.g. ironic -> bm) and fail hard if
        # underlying hardware type does not support setting boot mode as
        # it seems to be a hopeless misconfiguration
        _set_boot_mode_on_bm(task, ironic_boot_mode, fail_if_unsupported=True)

def is_secure_boot_requested(node):
    """Returns True if secure_boot is requested for deploy.

    This method checks node property for secure_boot and returns True
    if it is requested.

    :param node: a single Node.
    :raises: InvalidParameterValue if the capabilities string is not a
             dictionary or is malformed.
    :returns: True if secure_boot is requested.
    """

    capabilities = common_utils.parse_instance_info_capabilities(node)
    sec_boot = capabilities.get('secure_boot', 'false').lower()

    return sec_boot == 'true'

def get_boot_mode_for_deploy(node):
    """Returns the boot mode that would be used for deploy.

    This method returns boot mode to be used for deploy.
    It returns 'uefi' if 'secure_boot' is set to 'true' in
    'instance_info/capabilities' of node.
    Otherwise it returns value of 'boot_mode' in 'properties/capabilities'
    of node if set. If that is not set, it returns boot mode in
    'internal_driver_info/deploy_boot_mode' for the node.
    If that is not set, it returns boot mode in
    'instance_info/deploy_boot_mode' for the node.
    It would return None if boot mode is present neither in 'capabilities' of
    node 'properties' nor in node's 'internal_driver_info' nor in node's
    'instance_info' (which could also be None).

    :param node: an ironic node object.
    :returns: 'bios', 'uefi' or None
    :raises: InvalidParameterValue, if the node boot mode disagrees with
        the boot mode set to node properties/capabilities
    """

    if is_secure_boot_requested(node):
        LOG.debug('Deploy boot mode is uefi for %s.', node.uuid)
        return 'uefi'

    # NOTE(etingof):
    # The search for a boot mode should be in the priority order:
    #
    # 1) instance_info.capabilities
    # 2) instance_info.deploy_boot_mode (deprecated in Wallaby)
    # 3) properties.capabilities
    # 4) driver_internal_info.deploy_boot_mode (internal)
    #
    # Because:
    #
    # (1) and (2) are deleted during teardown
    # (4) will never be touched if node properties/capabilities
    #     are still present.
    # (3) becomes operational default as the last resort

    inst_boot_mode = (
        common_utils.parse_instance_info_capabilities(node).get('boot_mode')
    )
    cap_boot_mode = driver_utils.get_node_capability(node, 'boot_mode')

    old_boot_mode = node.instance_info.get('deploy_boot_mode')
    if old_boot_mode:
        LOG.warning('Using instance_info/deploy_boot_mode is deprecated, '
                    'please use instance_info/capabilities with boot mode '
                    'for node %s', node.uuid)

    boot_mode = (
        inst_boot_mode
        or old_boot_mode
        or cap_boot_mode
        or node.driver_internal_info.get('deploy_boot_mode')
    )

    if not boot_mode:
        return

    boot_mode = boot_mode.lower()

    # NOTE(etingof):
    # Make sure that the ultimate boot_mode agrees with the one set to
    # node properties/capabilities. This locks down node to use only
    # boot mode specified in properties/capabilities.
    # TODO(etingof): this logic will have to go away when we switch to traits
    if cap_boot_mode:
        cap_boot_mode = cap_boot_mode.lower()
        if cap_boot_mode != boot_mode:
            msg = (_("Node %(uuid)s boot mode %(boot_mode)s violates "
                     "node properties/capabilities %(caps)s") %
                   {'uuid': node.uuid,
                    'boot_mode': boot_mode,
                    'caps': cap_boot_mode})
            LOG.error(msg)
            raise exception.InvalidParameterValue(msg)

    LOG.debug('Deploy boot mode is %(boot_mode)s for %(node)s.',
              {'boot_mode': boot_mode, 'node': node.uuid})

    return boot_mode

def get_boot_mode(node):
    """Returns the boot mode.

    :param node: an ironic node object.
    :returns: 'bios' or 'uefi'
    :raises: InvalidParameterValue, if the node boot mode disagrees with
        the boot mode set to node properties/capabilities
    """
    boot_mode = get_boot_mode_for_deploy(node)
    if boot_mode:
        return boot_mode
    return CONF.deploy.default_boot_mode

@task_manager.require_exclusive_lock
def configure_secure_boot_if_needed(task):
    """Configures secure boot if it has been requested for the node."""
    if not is_secure_boot_requested(task.node):
        return

    try:
        task.driver.management.set_secure_boot_state(task, True)
    except exception.UnsupportedDriverExtension:
        # TODO(dtantsur): make a failure in Xena
        LOG.warning('Secure boot was requested for node %(node)s but its '
                    'management interface %(driver)s does not support it. '
                    'This warning will become an error in a future release.',
                    {'node': task.node.uuid,
                     'driver': task.node.get_interface('management')})
    except Exception as exc:
        with excutils.save_and_reraise_exception():
            LOG.error('Failed to configure secure boot for node %(node)s: '
                      '%(error)s',
                      {'node': task.node.uuid, 'error': exc},
                      exc_info=not isinstance(exc, exception.IronicException))
    else:
        LOG.info('Secure boot has been enabled for node %s', task.node.uuid)
        manager_utils.node_cache_boot_mode(task)

@task_manager.require_exclusive_lock
def deconfigure_secure_boot_if_needed(task):
    """Deconfigures secure boot if it has been requested for the node."""
    if not is_secure_boot_requested(task.node):
        return

    try:
        task.driver.management.set_secure_boot_state(task, False)
    except exception.UnsupportedDriverExtension:
        # NOTE(dtantsur): don't make it a hard failure to allow tearing down
        # misconfigured nodes.
        LOG.debug('Secure boot was requested for node %(node)s but its '
                  'management interface %(driver)s does not support it.',
                  {'node': task.node.uuid,
                   'driver': task.node.get_interface('management')})
    except Exception as exc:
        with excutils.save_and_reraise_exception():
            LOG.error('Failed to deconfigure secure boot for node %(node)s: '
                      '%(error)s',
                      {'node': task.node.uuid, 'error': exc},
                      exc_info=not isinstance(exc, exception.IronicException))
    else:
        LOG.info('Secure boot has been disabled for node %s', task.node.uuid)
        manager_utils.node_cache_boot_mode(task)


from django.db import models
from localflavor.br.br_states import STATE_CHOICES

class Campaign(models.Model):
    name = models.CharField('Nome da Campanha', max_length=150)
    description = models.TextField('Descrição da Campanha')
    date_created = models.DateField('Criado em', auto_now_add=True)
    last_updated = models.DateField('Última atualização em', auto_now=True)
    template_text = models.TextField('Template para telegrama, fax, email')

    class Meta:
        verbose_name = 'Campanha'

    def __str__(self):
        return self.name

class Argument(models.Model):
    campaign = models.ForeignKey('Campaign', verbose_name='Campanha', on_delete=models.CASCADE)
    text = models.TextField('Texto do argumento')

class Politician(models.Model):
    PARTY_CHOICES = (
        ('PMDB', 'PMDB - Partido Do Movimento Democrático Brasileiro'),
        ('PTB', 'PTB - Partido Trabalhista Brasileiro'),
        ('PDT', 'PDT - Partido Democrático Trabalhista'),
        ('PT', 'PT - Partido Dos Trabalhadores'),
        ('DEM', 'DEM - Democratas'),
        ('PCdoB', 'PCdoB - Partido Comunista Do Brasil'),
        ('PSB', 'PSB - Partido Socialista Brasileiro'),
        ('PSDB', 'PSDB - Partido Da Social Democracia Brasileira'),
        ('PTC', 'PTC - Partido Trabalhista Cristão'),
        ('PSC', 'PSC - Partido Social Cristão'),
        ('PMN', 'PMN - Partido Da Mobilização Nacional'),
        ('PRP', 'PRP - Partido Republicano Progressista'),
        ('PPS', 'PPS - Partido Popular Socialista'),
        ('PV',  'PV - Partido Verde'),
        ('PTdoB','PTdoB - Partido Trabalhista Do Brasil'),
        ('PP',  'PP - Partido Progressista'),
        ('PSTU', 'PSTU - Partido Socialista Dos Trabalhadores Unificado'),
        ('PCB', 'PCB - Partido Comunista Brasileiro'),
        ('PRTB', 'PRTB - Partido Renovador Trabalhista Brasileiro'),
        ('PHS', 'PHS - Partido Humanista Da Solidariedade'),
        ('PSDC', 'PSDC - Partido Social Democrata Cristão'),
        ('PCO', 'PCO - Partido Da Causa Operária'),
        ('PTN', 'PTN - Partido Trabalhista Nacional'),
        ('PSL', 'PSL - Partido Social Liberal'),
        ('PRB', 'PRB - Partido Republicano Brasileiro'),
        ('PSOL', 'PSOL - Partido Socialismo E Liberdade'),
        ('PR',  'PR - Partido Da República'),
        ('PSD', 'PSD - Partido Social Democrático'),
        ('PPL', 'PPL - Partido Pátria Livre'),
        ('PEN', 'PEN - Partido Ecológico Nacional'),
        ('PROS', 'PROS - Partido Republicano Da Ordem Social'),
        ('SDD ', 'SDD - Solidariedade'),
    )

    CATEGORY_CHOICES = (
        ('T', 'Titular'),
        ('S', 'Suplente'),
        ('E', 'Efetivo')
    )

    name = models.CharField('Nome', max_length=150)
    parliamentary_name = models.CharField('Nome Parlamentar', max_length=150)
    political_party = models.CharField('Partido', choices=PARTY_CHOICES, max_length=10)
    uf = models.CharField('UF', choices=STATE_CHOICES, max_length=10)
    category = models.CharField('Categoria', choices=CATEGORY_CHOICES, max_length=20)
    annex = models.IntegerField('Anexo', blank=True, null=True)
    chamber = models.IntegerField('Gabinete', blank=True, null=True)
    tel = models.CharField('Telefone', max_length=20)
    fax = models.CharField('Fax', max_length=20)
    email = models.EmailField('Email', max_length=100)
    date_created = models.DateField('Criado em', auto_now_add=True)
    last_updated = models.DateField('Última atualização em', auto_now=True)

    class Meta:
        verbose_name = 'Político'

    def __str__(self):
        return self.name

class Contact(models.Model):

    CONTACT_CHOICES = (
        ('tel', 'Telefone'),
        ('telegram', 'Telegrama via Web'),
        ('fax', 'Fax'),
        ('email', 'E-mail')
    )

    RESULT_CHOICES = (
        ('Telefonei', (
                ('10', 'Falei com o(a) Deputado(a) em pessoa.'),
                ('11', 'Falei com outra pessoa.'),
                ('12', 'Deixei recado em uma máquina.'),
                ('13', 'Fone: número ocupado'),
                ('14', 'Fone: ninguém atendeu.'),
                ('15', 'Fone: número inexistente ou outra falha.')
            )
        ),
        ('Enviei telegrama', (
                ('20', 'Nada a reportar: correio vai entregar no gabinete!'),
            )
        ),
        ('Enviei fax', (
                ('30', 'Fax: transmissão bem sucedida.'),
                ('31', 'Fax: número ocupado.'),
                ('32', 'Fax: não atendeu.'),
                ('33', 'Fax: número inexistente ou outra falha.')
            )
        ),
        ('Enviei e-mail', (
                ('40', 'E-mail enviado e não voltou, tomara que leiam.'),
                ('41', 'E-mail voltou com erro.'),
            )
        ),
    )

    politician = models.ForeignKey(Politician, verbose_name='Político', on_delete=models.CASCADE)
    campaign = models.ForeignKey(Campaign, verbose_name='Campanha', on_delete=models.CASCADE)
    contacted_by = models.CharField(verbose_name='Contato via', choices=CONTACT_CHOICES, max_length=10)
    result = models.CharField(verbose_name='Resultado do contato', choices=RESULT_CHOICES, max_length=10, blank=True, null=True)
    date_created = models.DateField('Criado em', auto_now_add=True)

    class Meta:
        verbose_name = 'Contato'

    def __str__(self):
        return '%s - %s' % (self.campaign.name, self.politician.name)

class HelpText(models.Model):
    campaign = models.ForeignKey(Campaign, verbose_name='Texto de ajuda para a campanha', on_delete=models.CASCADE)
    name = models.CharField('Nome da Ajuda', max_length=150)
    description = models.TextField('Texto de Ajuda')
    date_created = models.DateField('Criado em', auto_now_add=True)
    last_updated = models.DateField('Última atualização em', auto_now=True)

    class Meta:
        verbose_name = 'Texto de Ajuda'
        verbose_name_plural = 'Textos de Ajuda'

    def __str__(self):
        return self.name


import csv
from os import write
import itertools

w_filename = 'author_same_year.csv'
r_filename = 'EE226-2021spring-problem1/author_paper_all_with_year.csv'

with open(w_filename, 'w', newline='') as wf:
    writer = csv.writer(wf)
    writer.writerow(['authorid','authorid'])
    years = [[] for i in range(5)] #2016-2020
    with open(r_filename) as f:
        reader = csv.reader(f)
        head = next(reader)
        data = list(reader)
        for k,row in enumerate(data):
            print(k)
            if row[2] == '2016':
                if row[0] not in years[0]:
                    years[0].append(row[0])
            if row[2] == '2017':
                if row[0] not in years[1]:
                    years[1].append(row[0])
            if row[2] == '2018':
                if row[0] not in years[2]:
                    years[2].append(row[0])
            if row[2] == '2019':
                if row[0] not in years[3]:
                    years[3].append(row[0])
            if row[2] == '2020':
                if row[0] not in years[4]:
                    years[4].append(row[0])

        for year in years:
            print(years.index(year))
            print(len(year))
            tmp = list(itertools.permutations(year, 2))
            author_same_year = []
            for i in tmp:
                author_same_year.append(list(i))
            for i in author_same_year:
                writer.writerow(i)


import argparse
import asyncio
from contextlib import suppress
from pathlib import Path
from typing import Coroutine, List

import tomli

from aioraft import Raft
from aioraft.client import GrpcRaftClient
from aioraft.server import GrpcRaftServer
from aioraft.types import RaftState
from aioraft.utils import build_loopback_ip

_cleanup_coroutines: List[Coroutine] = []

def load_config():
    path = Path(__file__).parent / "config.toml"
    return tomli.loads(path.read_text())

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", "-p", type=int, default=50051)
    return parser.parse_args()

async def _main():
    args = parse_args()
    public_ip = build_loopback_ip()
    public_id = f"{public_ip}:{args.port}"

    config = load_config()
    configuration = tuple(
        server
        for server in config["raft"]["configuration"]
        if not server.endswith(str(args.port))
    )

    async def _on_state_changed(next_state: RaftState):
        print(f"[_on_state_changed] next_state: {next_state}")

    server = GrpcRaftServer()
    client = GrpcRaftClient()
    raft = await Raft.new(
        public_id,
        server=server,
        client=client,
        configuration=configuration,
        on_state_changed=_on_state_changed,
    )

    done, pending = await asyncio.wait(
        {
            asyncio.create_task(
                server.run(
                    host="0.0.0.0",
                    port=args.port,
                    cleanup_coroutines=_cleanup_coroutines,
                ),
            ),
            asyncio.create_task(raft.main()),
        },
        return_when=asyncio.FIRST_EXCEPTION,
    )
    for task in pending:
        task.cancel()
        with suppress(asyncio.CancelledError):
            await task

async def main():
    try:
        await _main()
    finally:
        await asyncio.gather(*_cleanup_coroutines)

if __name__ == "__main__":
    with suppress(KeyboardInterrupt):
        asyncio.run(main())


import sys
import os
import numpy as np
import pandas as pd
import copy
import itertools
import threading
from multiprocessing import Process
import datetime
import time
import sched

import edge
import globalVar
from globalVar import enumVar
from globalVar import enumScene
import util
import PkgAlg

class MyThread(threading.Thread):
    def __init__(self, threadLock, threadID, name, function, para):
        threading.Thread.__init__(self)
        self.threadLock = threadLock
        self.threadID = threadID
        self.name = name
        self.function = function
        self.para = para

    def run(self):
        # self.threadLock.acquire()
        print(f'thread {self.threadID} start')
        self.function(self.para)
        print(f'thread {self.threadID} stop')
        # self.threadLock.release()

def single_thread(switch, context):
    title_str = None
    title_list = None
    log_array = []
    f1_trace_array = []
    f2_trace_array = []

    #开始跑算法
    for i, query in enumerate(context.query_list):
        info_para = query.show_param()

        #对每个算法启动定时任务
        """
        print("job start")
        s = sched.scheduler(time.time, time.sleep)
        e1 = s.enter(0, 0, perform, ('Event1',)) #每10秒执行
        t = threading.Thread(target=s.run)
        t.start()
        """

        #运行算法
        util.sta_show_now_time_msg(msg=info_para)
        run_start = util.sta_get_start_time()
        switch[query.alg_name](query)
        run_seconds = util.sta_get_diff_time(run_start)
        info_result = query.show_result()
        util.sta_show_now_time_msg(msg=info_result+'\n')
        #记录算法的运行时间到csv文件
        #'data,alg,run,seconds'
        run_time_record=f'{query.folder_name}-{query.scene_idx+1},{query.alg_name},{query.run_idx+1},{run_seconds}\n'
        util.save_run_time(query.config.rum_time_path, run_time_record)

        #对每个算法停止定时任务
        """
        s.cancel(e1)
        t.join()
        print("job end")
        """

    #算法跑完后，开始统计数据
    for i, query in enumerate(context.query_list):
        # get title once in the 'for' section
        if (i == 0):
            title_str, title_list = query.show_title()

        log_str, log_list, f1_trace_list, f2_trace_list = query.show_log()
        log_array.append(log_list)
        #trace_list, f2_trace_list是一个算法分别对两个任务运行一次的最优值记录
        f1_trace_array.append(f1_trace_list)
        f2_trace_array.append(f2_trace_list)
    return title_str, title_list, log_array, f1_trace_array, f2_trace_array

def multi_thread(switch, context):
    title_str = None
    title_list = None
    log_array = []
    #f2_log_array = []

    threadLock = threading.Lock()
    threadID = 0
    threads = []

    for i, query in enumerate(context.query_list):
        threadID += 1
        t = MyThread(threadLock, threadID, query.alg_name, switch[query.alg_name], query)
        threads.append(t)
        # t.setDaemon(True)
        t.start()

    for t in threads:
        t.join()

    for i, query in enumerate(context.query_list):
        # get title once in the 'for' section
        if (i == 0):
            title_str, title_list = query.show_title()

        log_str, log_list, f1_trace_list, f2_trace_list = query.show_log()
        log_array.append(log_list)
        #f2_log_array.append(f2_log_list)
    return title_str, title_list, log_array #, f2_log_array

def multi_process(switch, context):
    title_str = None
    title_list = None
    log_array = []

    process_id = 0
    process_list = []

    for i, query in enumerate(context.query_list):
        process_id += 1
        p = Process(target=switch[query.alg_name], args=(query,))
        p.start()

    [p.join() for p in process_list]

    for i, query in enumerate(context.query_list):
        # get title once in the 'for' section
        if (i == 0):
            title_str, title_list = query.show_title()

        log_str, log_list, f1_trace_list, f2_trace_list = query.show_log()
        log_array.append(log_list)
    return title_str, title_list, log_array

"""
每隔10秒输出一次数据
"""
def perform(name):
    #获得当前fitness值
    nowTime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    start_time = time.strftime('%Y-%m-%d_%H.%M.%S',time.localtime(time.time()))
    print(f'{name} each 10 seconds: {nowTime}    {start_time}')

def begin_query(context):
    """ start query.

    Args:
        context: folder_name, slot_idx, query_list

    Returns:
        Query result class in globalVar.py
    """
    config = context.config
    switch = {
        'Single_GA': PkgAlg.run_single_GA,
        'Single_MFEA_II': PkgAlg.run_single_mfea_ii,
        'Single_MFEA_adaptive_cluster': PkgAlg.run_Single_MFEA_adaptive_cluster, # This is MFEA-AAT
        'Single_MFEA_adaptive_meme': PkgAlg.run_Single_MFEA_adaptive_meme,
        'Single_MFEA_adaptive_meme2': PkgAlg.run_Single_MFEA_adaptive_meme2,
        'Single_MFEA_adaptive_meme3': PkgAlg.run_Single_MFEA_adaptive_meme3,  #This is MFMA-AAT
    }

    '''
    start_time = util.sta_get_start_time()
    title_str, title_list, log_array = multi_thread(switch, context)
    util.sta_show_used_time_msg(start_time, msg='multi_thread')

    start_time = util.sta_get_start_time()
    title_str, title_list, log_array = multi_process(switch, context)
    util.sta_show_used_time_msg(start_time, msg='multi_process')
    input('pause')
    '''

    start_time = util.sta_get_start_time()
    title_str, title_list, log_array, f1_trace_array, f2_trace_array = single_thread(switch, context)
    util.sta_show_used_time_msg(start_time, msg='single_thread')

    table_all = np.array(log_array)
    f1_trace_all = np.array(f1_trace_array)
    #['SN1'-'1'-'Single_MFEA'-'10_1000_1000_0.9_0.1_0.1', '1',trace_array]
    #log_array只记录了每个算法的参数及最优值结果，trace_array记录整体过程
    #trace_array中的一行表示一次run中1*max_fit个依次找到的最优值数组
    #需要把trace_array转换成run*batch数据
    #len_row = len(trace_array)
    #len_col = len(trace_array[0])
    #trace_all = trace_all.reshape(len_row,len_col)
    f2_trace_all = np.array(f2_trace_array)
    # f1_trace_all和f2_trace_all中的每一行都是一个list，每个list是一个查询条件的记录
    #f2_trace_all = f2_trace_all.reshape(len_row,len_col)

    if 'Multi' in context.query_list[0].alg_name:
        sta_multi(context, table_all, title_list)
    else:
        # 单目标结果统计
        sta_single(context, table_all, title_list, f1_trace_all, f2_trace_all)

def sta_single(context, table_all, title_list, f1_trace_all, f2_trace_all):

    # after all query classes finish, begin statistic.
    """
    title_str = 'data_name,slot,alg,para,run,best_gen,best_fit,best_val'
    title_list = ['data_name', 'slot', 'alg', 'para', 'run',
                'best_gen', 'best_fit', 'best_val']
    """
    config = context.config
    file_name = context.folder_name + "_slot" + str(context.scene_idx + 1)
    edge.save_table_to_csv(context.config, table_all, title_list, file_name, prefix='sta', surfix='all')

    df = pd.DataFrame(table_all)
    df.columns = title_list

    # title_list = ['data_name', 'slot', 'alg', 'pop_size',
    #                'max_gen', 'max_fit', 'pc', 'pm', 'pl', 'run',
    #                'best_gen', 'best_fit', 'best_val']

    alg_list = df['alg'].unique()   # ['Single_MFEA' 'Single_mfea_akt']
    para_list = df['para'].unique()  # ['10_1000_1000_0.9_0.1_0.1']

    dict_gen = {}
    dict_fit = {}
    dict_val = {}
    combination = list(itertools.product(alg_list, para_list))
    for comb in combination:
        #print(comb): ('Single_MFEA', '10_1000_1000_0.9_0.1_0.1')
        col_key = ''.join(comb)
        pd_sub = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]
        pd_sub_gen = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_gen']
        pd_sub_fit = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_fit']
        pd_sub_val = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_val']

        dict_gen[col_key] = pd_sub_gen.values.astype(int)  # .tolist()  #.#T#.tolist()
        dict_fit[col_key] = pd_sub_fit.values.astype(int)  # .tolist()  #.T#.tolist()
        dict_val[col_key] = pd_sub_val.values.astype(float)  # .tolist()   #.apply(pd.to_numeric())  #.T#.tolist()

    file_path_list = []
    file_path_list.append(save_dict_to_csv(config, dict_gen, file_name, prefix='sta', surfix='gen'))
    file_path_list.append(save_dict_to_csv(config, dict_fit, file_name, prefix='sta', surfix='fit'))
    file_path_list.append(save_dict_to_csv(config, dict_val, file_name, prefix='sta', surfix='val'))
    util.computePvalue(file_path_list)

    #存储trace_all, f2_trace_all
    #获取算法数据名称参数
    #保存trace

    sta_trace(config, 'trace', f1_trace_all, f2_trace_all)
    print("a single task finish.")

def sta_trace(config, trace_name, trace_all_f1, trace_all_f2):
    #pre = trace_name.split('_')[0] # trace_f1
    #sur = trace_name.split('_')[1] # f1  or  f2
    pre = trace_name
    sur = ''

    #config.run_count

    trace_title_str=['data-slot-alg-para','run','trace']
    # 共三列，第一列是data-slot-alg-para，第二列是run_index，第三列是trace
    df_f1 = pd.DataFrame(trace_all_f1)
    df_f1.columns = trace_title_str

    df_f2 = pd.DataFrame(trace_all_f2)
    df_f2.columns = trace_title_str
    #print(df.shape[0])
    #print(df)
    file_list_f1 = df_f1['data-slot-alg-para'].unique()
    file_list_f2 = df_f2['data-slot-alg-para'].unique()
    #判断file_list_f1和file_list_f2是否相等
    file_list=[]
    file_list_f1_str = str(file_list_f1)
    file_list_f2_str = str(file_list_f2)
    if( file_list_f1_str == file_list_f2_str ):
        file_list = file_list_f1
    if len(file_list) == 0:
        print(f'file_list_f1\n{file_list_f1}\nfile_list_f2\n{file_list_f2}')
        input("check file_list error")
    #第一列是文件名，第二列是run_index，第三列是trace数据
    #按第一列和第二列排序
    df_sort_f1 = df_f1.sort_values(by=['data-slot-alg-para','run'], ascending=[True, True])
    df_sort_f2 = df_f2.sort_values(by=['data-slot-alg-para','run'], ascending=[True, True])
    #print(df_sort)
    for file in file_list:
        df_sort_sub_f1 = df_sort_f1[(df_sort_f1[u'data-slot-alg-para'] == file)]
        df_trace_f1 = df_sort_sub_f1['trace']

        df_sort_sub_f2 = df_sort_f2[(df_sort_f2[u'data-slot-alg-para'] == file)]
        df_trace_f2 = df_sort_sub_f2['trace']
        #df_trace_val = df_trace.values.astype(float)

        file_name = file

        edge.save_table_to_csv(config, df_trace_f1, ['trace_f1'],
            file_name, prefix=pre, surfix='f1', second_folder='trace')
        edge.save_table_to_csv(config, df_trace_f2, ['trace_f2'],
            file_name, prefix=pre, surfix='f2', second_folder='trace')

        second_folder="trace_pkl"
        #folderPath = config.outputPath+config.pathSep
        #path_test = folderPath+"record"+config.pathSep+second_folder+config.pathSep
        folderPath = globalVar.cfg.log_folder_path+globalVar.pathSep
        path_test = folderPath+config.pathSep+second_folder+config.pathSep
        if not os.path.exists(path_test):
            os.makedirs(path_test)

        path = path_test+file_name+'.pkl'

        # 保存pickle文件，同时包含trace_f1和df_trace_f2
        obj_trace_f1_f2 = (df_trace_f1, df_trace_f2)

        edge.save_trace_to_pickle(path, obj_trace_f1_f2)

def sta_multi(context, table_all, title_list):
    # after all query classes finish, begin statistic.
    config = context.config

    file_name = context.folder_name + "_slot" + str(context.scene_idx + 1)
    edge.save_table_to_csv(context.config, table_all, title_list, file_name, prefix='sta', surfix='all')

    df = pd.DataFrame(table_all)
    df.columns = title_list

    # title_list = ['data_name', 'slot', 'alg', 'pop_size',
    #                'max_gen', 'max_fit', 'pc', 'pm', 'pl', 'run',
    #                'best_gen', 'best_fit', 'best_f1', 'best_f2']

    alg_list = df['alg'].unique()
    para_list = df['para'].unique()

    dict_gen = {}
    dict_fit = {}
    dict_f1 = {}
    dict_f2 = {}
    combination = list(itertools.product(alg_list, para_list))
    for comb in combination:
        col_key = ''.join(comb)
        pd_sub = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]
        pd_sub_gen = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_gen']
        pd_sub_fit = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_fit']
        pd_sub_f1 = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_f1']
        pd_sub_f2 = df[(df[u'alg'] == comb[0]) & (df[u'para'] == comb[1])]['best_f2']

        dict_gen[col_key] = pd_sub_gen.values.astype(int)  # .tolist()  #.#T#.tolist()
        dict_fit[col_key] = pd_sub_fit.values.astype(int)  # .tolist()  #.T#.tolist()
        dict_f1[col_key] = pd_sub_f1.values.astype(float)  # .tolist()   #.apply(pd.to_numeric())  #.T#.tolist()
        dict_f2[col_key] = pd_sub_f2.values.astype(float)  # .tolist()   #.apply(pd.to_numeric())  #.T#.tolist()

    file_path_list = []
    file_path_list.append(save_dict_to_csv(config, dict_gen, file_name, prefix='sta', surfix='gen'))
    file_path_list.append(save_dict_to_csv(config, dict_fit, file_name, prefix='sta', surfix='fit'))
    file_path_list.append(save_dict_to_csv(config, dict_f1, file_name, prefix='sta', surfix='f1'))
    file_path_list.append(save_dict_to_csv(config, dict_f2, file_name, prefix='sta', surfix='f2'))
    util.computePvalue(file_path_list)
    print()

def save_dict_to_csv(config, dict_data, file_name, prefix='', surfix=''):
    table_data, title = util.convert_dict_to_table(dict_data)
    file_path = edge.save_table_to_csv(config, table_data, title, file_name, prefix, surfix)
    return file_path

def gen_query_list(config, folder_idx, folder_name, user_count, user_group, list_scene, scene_idx, scene):
    """ Generate query class with query parameters.

    Args:
        folder_idx: current data folder idx.
        folder_name: current data folder name.

    Returns:
        Query class list in globalVar.py
    """

    query_cls_list = []

    folder_list = config.listFolder
    alg_list = config.listMethod
    para_list = config.listEvoPara

    folder_count = config.folder_count
    scene_count = config.scene_count
    alg_count = config.alg_count
    para_count = config.para_count
    run_count = config.run_count

    query_count = config.query_count

    cell_id_min = config.cellIdMin
    cell_id_max = config.cellIdMax

    for alg_idx, alg_name in enumerate(alg_list):
        # globalVar.query_idx += 1
        for para_idx, para_str in enumerate(para_list):
            # globalVar.query_idx += 1
            for run_idx in np.arange(run_count):
                #globalVar.query_idx += 1
                config.query_idx += 1
                pop_size, max_gen, max_fit, pc, pm, pl, pa, cross_indi, mutate_indi = edge.setNowEvoPara(config, para_str)

                query_cls = globalVar.Query(config, config.query_idx, query_count,
                                            folder_idx, folder_count, folder_name,
                                            user_count, user_group, list_scene, cell_id_min, cell_id_max,
                                            scene_idx, scene_count, scene,
                                            alg_idx, alg_count, alg_name,
                                            para_idx, para_count, para_str,
                                            run_idx, run_count,
                                            pop_size, max_gen, max_fit, pc, pm, pl, pa,
                                            cross_indi, mutate_indi)
                query_cls_list.append(query_cls)
    return query_cls_list

def load_data(config, folder_idx, folder_name):
    """ Load 4 csv files from folder

    Args:
        folder: current data folder name.

    Returens:
        None.
    """

    folder = folder_name
    print("\n#####run data folder: {}".format(folder))
    config.curDataFolder = folder
    # config.outputPath+os.sep+
    data_folder_name = config.dataFolder
    sep = config.pathSep

    fileUser = data_folder_name + sep + folder + sep + "user.csv"
    fileServ = data_folder_name + sep + folder + sep + "server.csv"
    fileMove = data_folder_name + sep + folder + sep + "move.csv"
    fileHop = data_folder_name + sep + folder + sep + "hopMatrix.xlsx"

    if not os.path.exists(fileUser):
        print(f'{fileUser} not exist')
    if not os.path.exists(fileServ):
        print(f'{fileServ} not exist')
    if not os.path.exists(fileMove):
        print(f'{fileMove} not exist')
    if not os.path.exists(fileHop):
        print(f'{fileHop} not exist')

    dfUser = util.readCsv(fileUser)
    dfServ = util.readCsv(fileServ)
    npUser = np.array(dfUser)
    npServ = np.array(dfServ)
    dfMove = util.readCsv(fileMove)
    matrix = util.readExcel2Matrix(fileHop)
    config.dfUser = dfUser
    config.dfServ = dfServ
    config.npUser = npUser
    config.npServ = npServ
    config.cellMatrix = matrix

    listClsUser = edge.loadUser(config, dfUser)
    listClsServ = edge.loadServ(config, dfServ)
    #globalVar.set_value(enumVar.listClsUser, listClsUser)
    #globalVar.set_value(enumVar.listClsServ, listClsServ)
    config.listClsUser = listClsUser
    config.listClsServ = listClsServ

    # important operation!
    userCount = len(listClsUser)
    config.userCount = userCount
    #globalVar.set_value(enumVar.userCount, userCount)

    # print("userCount is {}".format(userCount))
    # input()
    if userCount == 0:
        print("input error, userCount is 0")
        input()

    servCount = len(listClsServ)
    config.servCount = servCount
    if servCount == 0:
        print("input error, servCount is 0")
        input()

    listSlot = edge.loadSlot(dfMove)
    # conductSingleObjective(listSlot)

    #listScene = globalVar.get_value(enumVar.listScene)
    #listScene = config.listScene
    config.listScene.clear()

    initialScene = edge.initialSceneFunction(config, listSlot[0])
    # print("initialScene:\n{}".format(initialScene))
    edge.recursiveScene(config, listSlot, 0, initialScene)  # return listScene

    return userCount, config.listScene

def run_data_folder(config, folder_idx, folder_name):
    """  Run each data folder.

    Args:
        folder_idx: current data folder idx.
        folder_name: current data folder name.

    Returns:
        None. Output files in result folder.
    """

    user_count, list_scene_ori = load_data(config, folder_idx, folder_name)
    # list_scene_ori = globalVar.get_value(enumVar.listScene)
    #只取前slotNum个scene
    if config.slotNum <= 0 or config.slotNum>=len(list_scene_ori):
        list_scene = copy.deepcopy(list_scene_ori)
    else:
        list_scene = copy.deepcopy(list_scene_ori[0:config.slotNum])

    # user_count = globalVar.get_value(enumVar.userCount)
    # index = 0
    scene_count = len(list_scene)
    config.list_scene = list_scene
    config.scene_count = scene_count

    folder_count = config.folder_count
    alg_count = config.alg_count
    para_count = config.para_count
    run_count = config.run_count

    # when get list_scene, compute the total query count.
    config.query_count = folder_count * scene_count * alg_count * para_count * run_count

    valStack = np.zeros((5, scene_count))
    valStack[0, :] = np.arange(1, scene_count + 1)
    phenStack = np.arange(1, user_count + 1, 1)
    for scene_idx, scene in enumerate(list_scene):
        #globalVar.set_value(enumVar.currentScene, scene)
        config.current_scene = scene
        config.currentScene = scene

        # user_group是k-means分组
        pkl_group = config.dataFolder + config.pathSep + folder_name + config.pathSep + "cluster"+config.pathSep+"group_"+str(scene_idx+1)+".pkl"
        if not os.path.exists(pkl_group):
            print(f'{pkl_group} not exist')
            input("pause")
        user_group = edge.cacheLoadGroup(pkl_group)

        query_list = gen_query_list(config, folder_idx, folder_name, user_count, user_group, list_scene, scene_idx, scene)

        context = globalVar.Context(config, folder_name, scene_idx, query_list)
        info = f'{folder_name}_{scene_idx}'
        start = util.sta_show_now_time_msg(msg='begin a query '+info)
        begin_query(context)
        util.sta_show_used_time_msg(start, msg='')

def init_global_var(cfg_file):
    """ Initialize global variables.
    :param
        cfg_file: config file path.
    :return
        Config class.
    """

    #给record文件加开始时间
    nowTime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    start_time = time.strftime('%Y-%m-%d_%H.%M.%S',time.localtime(time.time()))

    config = globalVar.Config(cfg_file, start_time)
    globalVar.cfg = config

    # record program start time to globalVar.runTime
    run_time = util.setRunTime()
    config.runTime = run_time

    # corss platform: windows "\\", linux "/"
    config.pathSep = os.path.sep

    sys.setrecursionlimit(10000)

    config.query_idx = -1

    # remain old code
    use_old_code = 0
    if use_old_code == 1:
        globalVar._init()

        globalVar.set_value(enumVar.listFolder, globalVar.cfg.listFolder)
        globalVar.set_value(enumVar.listMethod, globalVar.cfg.listMethod)
        globalVar.set_value(enumVar.listEvoPara, globalVar.cfg.listEvoPara)

        globalVar.set_value(enumVar.cellMatrix, "")
        globalVar.set_value(enumVar.currentSlot, "")
        globalVar.set_value(enumVar.listScene, [])

        globalVar.set_value(enumVar.userCount, 0)
        globalVar.set_value(enumVar.cellIdMin, 1)
        globalVar.set_value(enumVar.cellIdMax, 18)
        globalVar.set_value(enumVar.servIdMin, 1)
        globalVar.set_value(enumVar.servIdMax, 18)

        globalVar.set_value(enumVar.nowFitnessCount, 0)

        # scene
        globalVar.set_value(enumScene._userId, int(0))
        globalVar.set_value(enumScene._uInCell, int(1))
        globalVar.set_value(enumScene._userRequ, int(2))
        globalVar.set_value(enumScene._inServPflNum, int(3))
        globalVar.set_value(enumScene._inServCapa, int(4))
        globalVar.set_value(enumScene._uInCellLast, int(5))
        globalVar.set_value(enumScene._userRequLast, int(6))
        globalVar.set_value(enumScene._inServPflNumLast, int(7))
        globalVar.set_value(enumScene._inServCapaLast, int(8))
        globalVar.set_value(enumScene._pInServLast, int(9))
        globalVar.set_value(enumScene._userQueuLast, int(10))
        globalVar.set_value(enumScene._userQueu, int(11))
        globalVar.set_value(enumScene._pInServ, int(12))

    return config

def main(cfg_file=''):
    """ program entry. """

    # cfg_folder = 'cfg_folder'
    # cfg_file_list = os.listdir(cfg_folder)

    # initialize globalVar config class 'globalVar.cfg'
    #globalVar._init()

    # cfg_file = r'./globalCfg.yml'
    #cfg_file = 'globalCfg.yml'
    config = init_global_var(cfg_file)

    config.rum_time_path = config.log_folder_path+'run_time_record.csv'
    util.create_run_time(config.rum_time_path,'data,alg,run_idx,seconds\n')

    # run each data folder
    folder_list = config.listFolder
    for folder_idx, folder_name in enumerate(folder_list):
        # globalVar.query_idx += 1
        run_data_folder(config, folder_idx, folder_name)

if (__name__ == "__main__"):
    # command: python main.py globalCfg.yml
    len_argv = len(sys.argv)
    cfg_file_path = ''
    if len_argv == 2:
        cfg_file_path = sys.argv[1]
    #cfg_file_path="globalCfg_single_obj_sn2.yml"
    main(cfg_file = cfg_file_path)
    # 对每一个算法，生成一个trace文件，trace文件中包含data(benchmark) dict()，每一列是第几次run，每一行是一次fitness


# **************************************************************************************************
# Name:    io_functions.py
# Purpose: Defines input and output functions for this project
# Input:   none
# Output:  input and output functions
# Notes:   none
# **************************************************************************************************

# ********************************************* Imports ********************************************
# Standard Library
import numpy as np
import pathlib
import csv
import matplotlib as plot
# My Library
import definitions
import classes.sample_class as sample_class
import utils.math_functions as math_func

# *************************************** Variable Definitions *************************************
DATA_PATH = pathlib.Path(definitions.ROOT_DIR + "/data/")
IMG_PATH = pathlib.Path(definitions.ROOT_DIR + "/img/Shear_Cryst_XY_Fig/")
SRC_PATH = pathlib.Path(definitions.ROOT_DIR + "/src/")
TEST_PATH = pathlib.Path(definitions.ROOT_DIR + "/test/")

# *************************************** Function Definitions *************************************
def read_hkl_from_csv(filename):
    # **********************************************************************************************
    # Name:    read_hkl_from_csv
    # Purpose: function that reads the hkl vector list from csv and returns as list
    # Input:   filename (string) - filename of the data to load, csv must be in n x 3 format
    # Output:  hkl_list (n x 3 matrix) - matrix that holds all the hkl vectors from csv file
    # Notes:   none
    # **********************************************************************************************

    # initialize return list and path
    hkl_list = np.empty((0, 3))
    file_path = DATA_PATH / filename

    # open csv storing hkl vectors
    with open(file_path) as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        line_count = 0

        # iterate over all rows in csv and add to hkl_list
        for row in csv_reader:
            hkl_vec = np.array([float(row[0]), float(row[1]), float(row[2])])
            hkl_list = np.vstack((hkl_list, hkl_vec))
            line_count += 1
        print(f'Processed {line_count} hkl vectors.')
    return hkl_list

def read_grains_from_csv(filename):
    # **********************************************************************************************
    # Name:    read_grains_from_csv
    # Purpose: function that reads grain data from csv and returns as list of grains
    # Input:   filename (string) - filename of data to load, csv must in n x 13 format
    # Output:  grain_list (list) - matrix that holds n grains from csv file
    # Notes:   none
    # **********************************************************************************************

    # initialize return list and path
    grain_list = []
    file_path = DATA_PATH / filename

    # open csv storing hkl vectors
    with open(file_path) as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        line_count = 0

        temp_cell = sample_class.UnitCell(np.array([4, 4, 4, 90, 90, 90]))

        # iterate over all rows in csv and add to grain list
        for row in csv_reader:
            if line_count != 0:
                dim_vec = np.array([float(row[0]), float(row[1]), float(row[2])])
                com_vec = np.array([float(row[3]), float(row[4]), float(row[5])])
                quat_vec = np.array([float(row[6]), float(row[7]), float(row[8]), float(row[9])])
                orient_vec = np.array([float(row[10]), float(row[11]), float(row[1])])

                # normalize vectors
                quat_vec = math_func.normalize(quat_vec)
                orient_vec = math_func.normalize(orient_vec)

                # Grain (unitCell, dimension, COM, orientation)
                temp_grain = sample_class.Grain(temp_cell, dim_vec, com_vec, quat_vec)
                temp_grain.vector2quat(orient_vec)

                # add to list
                grain_list.append(temp_grain)
            line_count += 1
        print(f'Processed {line_count} grains.')
    return grain_list

def save_matplotlib_plot(save_plot, filename='fig_1', extension='.png', tight=False):
    # **********************************************************************************************
    # Name:    save_matplotlib_plot
    # Purpose: function saves matplotlib plot to img directory
    # Input:   save_plot (matplotlib object) - plot object to save as image
    #          filename (string) - string of filename to save to
    #          extension (string) - extension of file type to save to
    #          tight (bool) - option for saving just the plot area of plot
    # Output:  none
    # Notes:   none
    # **********************************************************************************************

    # combine filename and extension
    path = IMG_PATH
    path = path / (filename + extension)

    print(path)

    # check if plot is to be saved as tight and save accordingly
    if tight:
        save_plot.savefig(path, bbox_inches='tight', pad_inches=0)
    else:
        save_plot.savefig(path, pad_inches=0)

def save_matplotlib_ani(save_ani, filename='fig_1', extension='.gif', tight=False):
    # **********************************************************************************************
    # Name:    save_matplotlib_plot
    # Purpose: function saves matplotlib animation to img directory
    # Input:   save_ani (matplotlib object) - animation object to save
    #          filename (string) - string of filename to save to
    #          extension (string) - extension of file type to save to
    #          tight (bool) - option for saving just the plot area of plot
    # Output:  none
    # Notes:   none
    # **********************************************************************************************

    # combine filename and extension
    path = IMG_PATH
    path = path / (filename + extension)

    print(path)

    # check if plot is to be saved as tight and save accordingly
    if tight:
        save_ani.save(path, writer='imagemagick', fps=5, bbox_inches='tight')
    else:
        save_ani.save(path, writer='imagemagick', fps=5)

def read_image_to_array(image_path):
    # **********************************************************************************************
    # Name:    read_image_to_array
    # Purpose: function reads image to numpy array
    # Input:   image_path (string) - path to image for reading to array
    # Output:  image_array (n x m x 4) - array of image pixel locations and RGBA values
    # Notes:   none
    # **********************************************************************************************

    # import image
    print("Read Image")

def write_hkl_to_csv(hkl_list, filename):
    # **********************************************************************************************
    # Name:    write_hkl_to_csv
    # Purpose: function that reads the hkl vector list from csv and returns as list
    # Input:   hkl_list (3 x n) - matrix that holds all the hkl vectors from csv file
    #          filename (string) - filename of the data to save
    # Output:  csv file with hkl list stored
    # Notes:   none
    # **********************************************************************************************

    # initialize path
    file_path = DATA_PATH / filename + '.csv'

    # save to csv
    np.savetxt(file_path, hkl_list, delimiter=',')