import numpy as np
from scipy.sparse.linalg import cg
from scipy.sparse.linalg import LinearOperator
from joblib import Parallel, delayed

from .sampler import *

class Trainer:
	def __init__(self, mol_info, reg_list=(100,0.9,1e-4), cores=8):
		self.h1e, self.eri, self.N_up, self.N_down = mol_info
		self.reg_list = reg_list
		self.step_count = 0
		self.nvar = 0
		self.parallel_cores = cores

	def train(self,wf,ecore,batch_size=100,num_steps=1000,init_state=None,\
		print_freq=20,state_file=None,param_file=None,out_freq=20):
		state = init_state
		# Energy list
		elist = np.zeros(num_steps+1,dtype=np.complex128)
		# early convergence
		convergence = 0
		for step in range(num_steps+1):
			updates, state, elist[step] = self.update_vector(wf,state,batch_size,gamma_fun(step),step,state_file)
			elist[step] += ecore
			self.apply_update(updates,wf)

			if (step!=0) and (abs(elist[step]-elist[step-1]) < 1e-3):
				convergence += 1
				if convergence > 10:
					print('Completed training step {}'.format(step))
					print("Current energy: {}".format(elist[step]))
					print("Energy converged! Program ends early!")
					break
			else:
				convergence = 0

			if (step % print_freq == 0) and (print_freq > 0):
				print('Completed training step {}'.format(step))
				print("Current energy: {}".format(elist[step]))

			if (out_freq>0) and (step%out_freq==0):
				wf.save_parameters(param_file+'Batch{}_Step{}'.format(batch_size,step))

		return wf, elist

	def eval(self,wf,ecore,batch_size=20000,init_state=None,param_file=None):
		if param_file != None:
			wf.load_parameters(param_file)
		if init_state == None:
			samp = Sampler(wf,self.N_up,self.N_down)
			samp.init_state('low')
		else:
			wf.init_table(init_state)
			samp = Sampler(wf,self.N_up,self.N_down)

		results = Parallel(n_jobs=self.parallel_cores)(\
			delayed(get_E)(samp,self) for batch in range(batch_size))

		return np.mean(results)+ecore

	def apply_update(self,updates,wf):
		wf.a += updates[0:wf.In]
		wf.b += updates[wf.In:wf.In+wf.Hi]
		wf.W += np.reshape(updates[wf.In+wf.Hi:],wf.W.shape)

	def update_vector(self,wf,init_state,batch_size,gamma,step,state_file=None,therm=True):
		self.nvar = wf.In + wf.Hi + wf.In*wf.Hi
		if init_state == None:
			samp = Sampler(wf,self.N_up,self.N_down)
			samp.init_state()
		else:
			wf.init_table(init_state)
			samp = Sampler(wf,self.N_up,self.N_down)

		results = Parallel(n_jobs=self.parallel_cores)(\
			delayed(get_sampler)(samp,self) for batch in range(batch_size))

		if therm == True:
			samp.thermalize(filename=state_file)

		elocals = np.array([i[0] for i in results])
		deriv_vectors = np.array([i[1] for i in results])

		forces = self.get_forces(elocals,deriv_vectors)
		Sv_operator = LinearOperator((self.nvar,self.nvar),dtype=np.complex128,\
			matvec = lambda v : self.cov_Sv(v,deriv_vectors,step))
		vec, info = cg(Sv_operator,forces)
		updates = -gamma * vec

		return updates, samp.state, np.mean(elocals)

	def get_forces(self,elocals,deriv_vectors):
		emean = np.mean(elocals)
		omean = np.mean(deriv_vectors.conj(),axis=0)
		correlator = np.mean(elocals.reshape(-1,1) * deriv_vectors.conj(),axis=0)

		return correlator - emean*omean

	def cov_Sv(self,v,deriv_vectors,step):
		# v(nvar,)
		# O = deriv_vectors(batch,nvar)
		# Ov(batch,)
		Ov = np.dot(deriv_vectors,v)
		# <O*,O>v (nvar,)
		term1 = np.dot(deriv_vectors.T.conj(), Ov) / deriv_vectors.shape[0]
		# <O*><O>v (nvar,)
		term2 = np.mean(deriv_vectors.conj(),axis=0) * np.mean(Ov)
		# reg(nvar,)
		reg = max(self.reg_list[0] * self.reg_list[1] ** step, self.reg_list[2]) * \
		        (np.mean(deriv_vectors.conj()*deriv_vectors,axis=0) - np.mean(deriv_vectors.conj(),axis=0)*\
		        	np.mean(deriv_vectors,axis=0)) * v

		return term1 - term2  + reg

	def get_deriv_vector(self,wf):
		nvar = wf.In + wf.Hi + wf.In*wf.Hi
		vector = np.zeros(nvar,dtype=np.complex128)
		vector[:wf.In] = wf.state[0]
		tanh_theta = np.tanh(wf.theta)
		vector[wf.In:wf.In+wf.Hi] = tanh_theta
		vector[wf.In+wf.Hi:] = np.dot(wf.state[0].reshape(-1,1), tanh_theta.reshape(1,-1)).ravel()

		return vector

	def get_elocal(self,wf):
		N_orbitals = wf.N_orbitals
		a = wf.a
		W = wf.W
		logpsi_2 = wf.logpsi_2
		theta = wf.theta

		def prob_1(excitation):
			occ = excitation[0] * N_orbitals + excitation[1]
			vir = excitation[0] * N_orbitals + excitation[2]
			logprob = 0 + 0j
			logprob += 2*a[vir] - 2*a[occ]
			logprob += np.sum(np.log(2*np.cosh( (theta-2*W[occ]+2*W[vir]) )) - logpsi_2)

			return np.exp(logprob)

		def prob_2(excitation):
			if excitation[0] == 2:
				# (up,down)
				occ1 = excitation[1][0]
				vir1 = excitation[2][0]
				occ2 = N_orbitals + excitation[1][1]
				vir2 = N_orbitals + excitation[2][1]
			else:
				occ1 = excitation[0] * N_orbitals + excitation[1][0]
				vir1 = excitation[0] * N_orbitals + excitation[2][0]
				occ2 = excitation[0] * N_orbitals + excitation[1][1]
				vir2 = excitation[0] * N_orbitals + excitation[2][1]
			logprob = 0 + 0j
			logprob += 2*a[vir1] + 2*a[vir2] - 2*a[occ1] - 2*a[occ2]
			logprob += np.sum(np.log(2*np.cosh( (theta - 2*W[occ1] - 2*W[occ2]	+ 2*W[vir1] + 2*W[vir2]) ))-logpsi_2)

			return np.exp(logprob)

		def sign_1(occ_orb,inx,occ,vir):
			step = 0
			if occ < vir:
				for i in range(inx+1,len(occ_orb)):
					if occ_orb[i] < vir:
						step += 1
					else:
						break
			else:
				for i in range(inx-1,-1,-1):
					if occ_orb[i] > vir:
						step += 1
					else:
						break

			return (-1)**step

		def sign_2(occ_orb,inx1,inx2,occ1,occ2,vir1,vir2):
			step = 0
			# o1,o2,v1,v2
			if vir1 > occ2:
				step += (inx2 - inx1 - 1)
				for i in range(inx2+1,len(occ_orb)):
					if occ_orb[i] > vir2:
						break
					elif occ_orb[i] < vir1:
						pass
					else:
						step += 1

			# v1,v2,o1,o2
			elif vir2 < occ1:
				step += (inx2 - inx1 - 1)
				for i in range(inx1-1,-1,-1):
					if occ_orb[i] < vir1:
						break
					elif occ_orb[i] > vir2:
						pass
					else:
						step += 1

			# v1,o1,o2,v2
			elif vir1<occ1 and occ2<vir2:
				for i in range(inx1-1,-1,-1):
					if occ_orb[i] < vir1:
						break
					else:
						step += 1
				for j in range(inx2+1,len(occ_orb)):
					if occ_orb[j] > vir2:
						break
					else:
						step +=1

			# o1,v1,o2,v2
			elif occ1<vir1 and occ2<vir2:
				for i in range(inx1+1,inx2):
					if occ_orb[i] > vir1:
						break
					else:
						step += 1
				for j in range(inx2+1,len(occ_orb)):
					if occ_orb[j] > vir2:
						break
					else:
						step +=1

			# v1,o1,v2,o2
			elif vir1<occ1 and vir2<occ2:
				for i in range(inx1-1,-1,-1):
					if occ_orb[i] < vir1:
						break
					else:
						step += 1
				for j in range(inx2-1,inx1,-1):
					if occ_orb[j] < vir2:
						break
					else:
						step += 1

			# o1,v1,v2,o2
			else:
				for i in range(inx1+1,inx2):
					if occ_orb[i] > vir1:
						break
					else:
						step += 1
				for j in range(inx2-1,inx1,-1):
					if occ_orb[j] < vir2:
						break
					else:
						step += 1

			return (-1)**step

		# now, start calsulating E_loc
		e_loc = 0.0 + 0.0j
		occ_1 = np.sort(wf.state[1]['occ_up'])
		occ_2 = np.sort(wf.state[1]['occ_down'])
		vir_1 = np.sort(wf.state[1]['vir_up'])
		vir_2 = np.sort(wf.state[1]['vir_down'])
		n_occ1 = len(occ_1)
		n_vir1 = len(vir_1)
		n_occ2 = len(occ_2)
		n_vir2 = len(vir_2)

		h1e = self.h1e
		eri = self.eri

		#<sigma|H|sigma>
		o1 = np.sum(h1e[occ_1,occ_1]) + np.sum(h1e[occ_2,occ_2])
		inx_1,inx_2 = np.meshgrid(occ_1,occ_1,indexing='ij')
		o2 = 0.5*np.sum(eri[inx_1,inx_2,inx_1,inx_2] - eri[inx_1,inx_2,inx_2,inx_1])
		inx_1,inx_2 = np.meshgrid(occ_2,occ_2,indexing='ij')
		o2 += 0.5*np.sum(eri[inx_1,inx_2,inx_1,inx_2] - eri[inx_1,inx_2,inx_2,inx_1])
		inx_1,inx_2 = np.meshgrid(occ_1,occ_2,indexing='ij')
		o2 += np.sum(eri[inx_1,inx_2,inx_1,inx_2])
		e_loc += o1+o2

		#<sigma|H|sigma_s>
		inx1,inx2,inx3,inx4 = np.meshgrid(occ_1,occ_1,vir_1,occ_1,indexing='ij')
		up_o1v1 = np.einsum('ipjp->ij',eri[inx1,inx2,inx3,inx4])
		inx1,inx2,inx3,inx4 = np.meshgrid(occ_1,occ_1,occ_1,vir_1,indexing='ij')
		up_o11v = np.einsum('ippj->ij',eri[inx1,inx2,inx3,inx4])
		inx1,inx2,inx3,inx4 = np.meshgrid(occ_1,occ_2,vir_1,occ_2,indexing='ij')
		up_o2v2 = np.einsum('ipjp->ij',eri[inx1,inx2,inx3,inx4])
		inx_1,inx_2 = np.meshgrid(occ_1,vir_1,indexing='ij')
		ov_1 = h1e[inx_1,inx_2] + up_o1v1 - up_o11v + up_o2v2
		for i in range(n_occ1):
			for j in range(n_vir1):
				occ = occ_1[i]
				vir = vir_1[j]
				e_loc += prob_1((0,occ,vir)) * sign_1(occ_1,i,occ,vir) * ov_1[i,j]
		up_o1v1,up_o11v,up_o2v2,ov_1 = (None,)*4

		inx1,inx2,inx3,inx4 = np.meshgrid(occ_2,occ_2,vir_2,occ_2,indexing='ij')
		down_o2v2 = np.einsum('ipjp->ij',eri[inx1,inx2,inx3,inx4])
		inx1,inx2,inx3,inx4 = np.meshgrid(occ_2,occ_2,occ_2,vir_2,indexing='ij')
		down_o22v = np.einsum('ippj->ij',eri[inx1,inx2,inx3,inx4])
		inx1,inx2,inx3,inx4 = np.meshgrid(occ_2,occ_1,vir_2,occ_1,indexing='ij')
		down_o1v1 = np.einsum('ipjp->ij',eri[inx1,inx2,inx3,inx4])
		inx_1,inx_2 = np.meshgrid(occ_2,vir_2,indexing='ij')
		ov_2 = h1e[inx_1,inx_2] + down_o2v2 - down_o22v + down_o1v1
		inx1,inx2,inx3,inx4,inx_1,inx_2 = (None,)*6
		for i in range(n_occ2):
			for j in range(n_vir2):
				occ = occ_2[i]
				vir = vir_2[j]
				e_loc += prob_1((1,occ,vir)) * sign_1(occ_2,i,occ,vir) * ov_2[i,j]
		down_o2v2,down_o22v,down_o1v1,ov_2 = (None,)*4

		#<sigma|H|sigma_d>
		for i in range(n_occ1):
			for j in range(i+1,n_occ1):
				for k in range(n_vir1):
					for l in range(k+1,n_vir1):
						occ1 = occ_1[i]
						occ2 = occ_1[j]
						vir1 = vir_1[k]
						vir2 = vir_1[l]
						o2 = eri[occ1,occ2,vir1,vir2] - eri[occ1,occ2,vir2,vir1]
						e_loc += prob_2((0,(occ1,occ2),(vir1,vir2))) * sign_2(occ_1,i,j,occ1,occ2,vir1,vir2) * o2

		for i in range(n_occ2):
			for j in range(i+1,n_occ2):
				for k in range(n_vir2):
					for l in range(k+1,n_vir2):
						occ1 = occ_2[i]
						occ2 = occ_2[j]
						vir1 = vir_2[k]
						vir2 = vir_2[l]
						o2 = eri[occ1,occ2,vir1,vir2] - eri[occ1,occ2,vir2,vir1]
						e_loc += prob_2((1,(occ1,occ2),(vir1,vir2))) * sign_2(occ_2,i,j,occ1,occ2,vir1,vir2) * o2

		for i in range(n_occ1):
			for j in range(n_occ2):
				for k in range(n_vir1):
					for l in range(n_vir2):
						occ1 = occ_1[i]
						occ2 = occ_2[j]
						vir1 = vir_1[k]
						vir2 = vir_2[l]
						o2 = eri[occ1,occ2,vir1,vir2]
						e_loc += prob_2((2,(occ1,occ2),(vir1,vir2))) * sign_1(occ_1,i,occ1,vir1) \
						 * sign_1(occ_2,j,occ2,vir2) * o2

		return e_loc

def gamma_fun(step):
	if step < 200:
		return 0.001
	elif step < 3000:
		return 0.001
	elif step < 5000:
		return 0.001
	else:
		return 1e-3

def get_sampler(mysampler,mytrainer):
	mysampler.thermalize()
	return mytrainer.get_elocal(mysampler.wf), mytrainer.get_deriv_vector(mysampler.wf)

def get_E(mysampler,mytrainer):
	mysampler.thermalize()
	return mytrainer.get_elocal(mysampler.wf)


import sys
import time
from PyQt5 import QtCore, QtGui, QtWidgets
from lib import vk_py_renderer

sys.path.append('lib\QtProperty')
sys.path.append('lib\libqt5')

from pyqtcore import QList
from qtvariantproperty import QtVariantEditorFactory, QtVariantPropertyManager
from qttreepropertybrowser import QtTreePropertyBrowser

from qtpropertymanager import (
    QtBoolPropertyManager,
    QtDoublePropertyManager,
    QtStringPropertyManager,
    QtColorPropertyManager,
    QtFontPropertyManager,
    QtPointPropertyManager,
    QtSizePropertyManager
)
from qteditorfactory import (
    QtDoubleSpinBoxFactory,
    QtCheckBoxFactory,
    QtSpinBoxFactory,
    QtLineEditFactory,
    QtEnumEditorFactory
)

import style

class VulkanWindow(QtGui.QWindow):

    def __init__(self, parent=None):
        super(VulkanWindow, self).__init__()
        self.vk_renderer = vk_py_renderer.VulkanRenderer()
        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self.render)

        self.fps_timer = time.perf_counter() * 1000
        self.last_elapsed_time =0
        self.fps = 0

    def __del__(self):
        cleanup()

    def initialize(self):
        self.vk_renderer.initialize(self.winId(), self.width(), self.height())
        self.timer.start()

    def render(self):
        self.vk_renderer.render()
        self.vk_renderer.update(self.last_elapsed_time)

        self.fps += 1
        if (time.perf_counter() * 1000 - self.fps_timer >= 1000):
            self.fps_timer = time.perf_counter() * 1000
            #print("FPS: %s" % self.fps)
            self.fps = 0
        self.last_elapsed_time = time.perf_counter()

    def cleanup(self):
        self.vk_renderer.cleanup()

    def resizeEvent(self, event):
        self.vk_renderer.resize(self.width(), self.height())

class MainWindow(QtWidgets.QMainWindow):

    def __init__(self):
        super(MainWindow, self).__init__()

        self.setGeometry(0, 0, 1280, 780)

        # Create actions and menus
        self.createMenu()
        self.createToolBar()

        # Create status bar
        self.statusBar().showMessage("Ready")

        # Vulkan widget
        self.vulkanWindow = VulkanWindow(self)
        self.vulkanWindowWidget = QtWidgets.QWidget.createWindowContainer(self.vulkanWindow)
        self.vulkanWindow.initialize()

        # Stop rendering while moving the window
        self.moveEventTimer = QtCore.QTimer(self)
        self.moveEventTimer.timeout.connect(self.moveEventDone)

        # Dock layout
        dock = QtWidgets.QDockWidget()
        dock.setAllowedAreas(QtCore.Qt.LeftDockWidgetArea | QtCore.Qt.RightDockWidgetArea)
        self.addDockWidget(QtCore.Qt.LeftDockWidgetArea, dock)

        # Property editor
        self.variantManager = QtVariantPropertyManager(self)
        self.variantManager.valueChangedSignal.connect(self.valueChanged)
        topItem = self.variantManager.addProperty(QtVariantPropertyManager.groupTypeId(), "Group Property")
        item = self.variantManager.addProperty(QtCore.QVariant.Bool, "Pause?")
        item.setValue(False)
        topItem.addSubProperty(item)
        self.variantFactory = QtVariantEditorFactory()
        self.propertyEditor = QtTreePropertyBrowser(dock)
        self.propertyEditor.setFactoryForManager(self.variantManager, self.variantFactory)
        self.propertyEditor.addProperty(topItem)
        self.propertyEditor.setPropertiesWithoutValueMarked(True)
        self.propertyEditor.setRootIsDecorated(False)

        dock.setWidget(self.propertyEditor)

        # Vulkan viewport
        dock = QtWidgets.QDockWidget()
        dock.setWidget(self.vulkanWindowWidget)

        # Box layout
        #layout = QtWidgets.QVBoxLayout()
        #layout.addWidget(self.button)
        #layout.addWidget(self.vulkanWindowWidget)
        #layout.addWidget(dock)

        #mainWidget = QtWidgets.QWidget()
        #mainWidget.setLayout(layout)

        #self.setCentralWidget(mainWidget)
        self.setCentralWidget(self.vulkanWindowWidget)

        #self.showMaximized()

    def __del__(self):
        del self.variantManager
        del self.variantFactory
        del self.variantEditor
        self.destroy()

    def valueChanged(self, property, value):
        tp = type(value)
        if tp == bool:
             self.vulkanWindow.vk_renderer.is_paused = value

    def closeEvent(self, event):
        self.vulkanWindow.cleanup()
        del self.vulkanWindow

    def moveEvent(self, event):
        self.moveEventTimer.start(500)
        self.vulkanWindow.timer.stop()

    def moveEventDone(self):
        self.vulkanWindow.timer.start()

    def createMenu(self):
        fileMenu = self.menuBar().addMenu("&File")
        fileMenu.addAction("Exit", self.exit)

    def createToolBar(self):
        exitAction = QtWidgets.QAction(QtGui.QIcon(":/images/open.png"), "&Open...", self)
        exitAction.setStatusTip("Open an existing file")
        exitAction.triggered.connect(self.exit)
        self.fileToolBar = self.addToolBar("File")
        self.fileToolBar.addAction(exitAction)

    def exit(self):
        print("exit")

if __name__ == '__main__':

    import sys
    app = QtWidgets.QApplication(sys.argv)
    app.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)

    style.set_app_style(app)

    win = MainWindow()
    win.show()

    def cleanup():
        print("exit")
        global win
        del win

    app.aboutToQuit.connect(cleanup)
    sys.exit(app.exec_())


import gym
from gym import spaces
import airsim
import numpy as np
import time
import math

class CarGym(gym.Env):
    def __init__(self):
        self.airsim_client = airsim.CarClient()
        self.airsim_client.confirmConnection()
        self.airsim_client.enableApiControl(True)
        print("API Control enabled: %s" % self.airsim_client.isApiControlEnabled())
        self.car_controls = airsim.CarControls()
        self.new_break_val = 0
        self.airsim_client.reset()

        self.state = {"position": np.zeros(3), "collision": None, "prev_position": np.zeros(3), "pose": None,
                      "prev_pose": None}

        # CarControls are:
        # throttle  0.0
        # steering = 0.0
        # brake = 0.0
        # handbrake = False
        # is_manual_gear = False
        # manual_gear = 0
        # gear_immediate = True
        # so action is the 3-element vector which increases/decreases ONLY throttle, steering, and brake.
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(3,), dtype=np.float32)
        self.observation_space = spaces.Box( low=0, high=255, shape=(1, 84, 84))

        # create an image request from camera "0", image type is depth perspective, pixels as float = True, compress= False
        self.image_request = airsim.ImageRequest(
            "0", airsim.ImageType.DepthPerspective, True, False
        )

    def _do_action(self, action):
        """
        transfoms action into carControls and feeds it into the airsim client(). Then it goes to sleep.
        :param action: a 7-element array, where each element increases/decreases the corresponding carControl() object
        :return: None
        """
        # throttle  0.0
        new_throttle = self.car_controls.throttle + action[0]
        if new_throttle >= 1.0:
            new_throttle = 1.0
        elif new_throttle < -1.0:
            new_throttle = -1.0
        self.car_controls.throttle = new_throttle

        # steering = 0.0
        new_steering = self.car_controls.steering + action[1]
        if new_steering >= 1.0:
            new_steering = 1.0
        elif new_steering <= -1.0:
            new_steering = -1.0
        self.car_controls.steering = new_steering

        # brake = 0.0
        if action[2] >= 0.5:
            self.car_controls.brake = True
        else:
            self.car_controls.brake = False
        car_controls = airsim.CarControls()
        car_controls.throttle = self.car_controls.throttle
        car_controls.steering = self.car_controls.steering
        # car_controls.brake = self.car_controls.brake

        self.airsim_client.setCarControls(car_controls)

    def transform_obs(self, responses):
        img1d = np.array(responses[0].image_data_float, dtype=np.float)
        #img1d is an ndarray of size (36864,). min=1.986 max=65504.0
        img1d = 255 / np.maximum(np.ones(img1d.size), img1d)
        #img1d is now an ndaray of size (36864,) min=0.00389 max=128.377
        img2d = np.reshape(img1d, (responses[0].height, responses[0].width))
        #img2d is ndarray(144,256)
        from PIL import Image

        image = Image.fromarray(img2d)
        im_final = np.array(image.resize((84, 84)).convert("L"))
        # im_final is ndarray(84,84)
        # image is now resized and converted to gray scale.
        return_im = im_final.reshape([1, 84, 84])
        # return_im is #ndarray(84,84,1)
        return return_im

    def _get_obs(self):
        responses = self.airsim_client.simGetImages([self.image_request])
        image = self.transform_obs(responses)

        self.car_state = self.airsim_client.getCarState()
        collision = self.airsim_client.simGetCollisionInfo().has_collided

        self.state["prev_pose"] = self.state["pose"]
        self.state["pose"] = self.car_state.kinematics_estimated
        self.state["collision"] = collision
        return image

    def _compute_reward(self):
        """
         reward is computed as reward_dist + reward_speed.
                reward_dist is either:
                    (-3), when the car's smallest distance to a set of points is bigger than a threshold.
                    (math.exp( -3 * dist) - 0.5) (if the dist is large, reward is small) otherwise
                reward_speed is:
                    (-) if the car's speed is below halfway of MIN_SPEED and MAX_SPEED
                    (0) if the car's speed is exactly halfway
                    (+) if the car's speed is above halfway
        on the other hand, the episode is done if:
                the reward is < 1
                brake's are off, but speed < 1
                there's a collision
        :return: reward, done
        """
        MAX_SPEED = 300
        MIN_SPEED = 10
        thresh_dist = 3.5
        beta = 3

        z = 0
        pts = [
            np.array([0, -1, z]),
            np.array([130, -1, z]),
            np.array([130, 125, z]),
            np.array([0, 125, z]),
            np.array([0, -1, z]),
            np.array([130, -1, z]),
            np.array([130, -128, z]),
            np.array([0, -128, z]),
            np.array([0, -1, z]),
        ]
        pd = self.state["pose"].position
        car_pt = np.array([pd.x_val, pd.y_val, pd.z_val])

        dist = 10000000
        for i in range(0, len(pts) - 1):
            dist = min(
                dist,
                np.linalg.norm(np.cross((car_pt - pts[i]), (car_pt - pts[i + 1])))
                / np.linalg.norm(pts[i] - pts[i + 1]),
                )

        # print(dist)
        if dist > thresh_dist:
            reward = -3
        else:
            reward_dist = math.exp(-beta * dist) - 0.5

            reward_speed = (
                                   (self.car_state.speed - MIN_SPEED) / (MAX_SPEED - MIN_SPEED)
                           ) - 0.5
            reward = reward_dist + reward_speed

        done = 0
        if reward < -1:
            print("Done because reward <-1")
            done = 1
        if self.car_controls.brake == 0:
            if self.car_state.speed <= 1:
                done = 1
        if self.state["collision"]:
            done = 1

        return reward, done

    def compute_distance(self, pos, target):
        xdiff_prev = (target["x_val"] - pos.x_val)
        xdiff_prev_sq = xdiff_prev ** 2
        ydiff_prev = (target["y_val"] - pos.y_val)
        ydiff_prev_sq = ydiff_prev ** 2
        zdiff_prev = (target["z_val"] - pos.z_val)
        zdiff_prev_sq = zdiff_prev ** 2
        _distance = math.sqrt(xdiff_prev_sq + ydiff_prev_sq + zdiff_prev_sq)

        return _distance

    def _compute__reward(self):
        """Reward #1: Customized reward system for getting to the end of the road.

        :return: reward, done
        """
        done = False
        reward = 0
        info = {}

        # target position
        target_pos = {
            "x_val": 101.336,
            "y_val": 0.077,
            "z_val": -0.587

        }
        # Get the car's current position and previous position
        curr_pos = self.state["pose"].position
        prev_pos = self.state["prev_pose"].position

        prev_distance = self.compute_distance(prev_pos, target_pos)
        curr_distance = self.compute_distance(curr_pos, target_pos)

        gained_distance = prev_distance - curr_distance
        reward = gained_distance
        info["reward"] = reward

        # check if done.
        if self.state["collision"]:
            reward = -100
            info["collision"] = True
            done = True

        return reward, done, info

    def step(self, action):
        self._do_action(action)
        obs = self._get_obs()
        reward, done, info = self._compute__reward()

        return obs, reward, done, self.state

    def reset(self):
        self.airsim_client.reset()
        return self._get_obs()

    def render(self, mode='human'):
        pass


import logging
import mock
import unittest
from datetime import datetime, timedelta
from logging import LogRecord
from mock import patch
from mock import MagicMock, PropertyMock

from google.api.monitored_resource_pb2 import MonitoredResource
from google.cloud.logging_v2.proto.log_entry_pb2 import LogEntryOperation
from google.protobuf import json_format as gcp_json_format, any_pb2 as gcp_any_pb2
import pytz

from curation_logging import curation_gae_handler
from curation_logging.curation_gae_handler import GCPStackDriverLogger, LogCompletionStatusEnum
from curation_logging.curation_gae_handler import GAE_LOGGING_MODULE_ID, GAE_LOGGING_VERSION_ID

LOG_BUFFER_SIZE = 3
SEVERITY_DEBUG = 100  # 100 is the equivalence of logging.DEBUG
SEVERITY_INFO = 200  # 200 is the equivalence of logging.INFO
SEVERITY_ERROR = 300  # 300 is the equivalence of logging.ERROR

class GCPStackDriverLoggerTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        print('**************************************************************')
        print(cls.__name__)
        print('**************************************************************')

    def setUp(self):
        self.project_id = 'aou-res-curation-test'
        self.request_method = 'GET'
        self.request_full_path = '/admin/v1/RemoveExpiredServiceAccountKeys'
        self.request_user_agent = 'AppEngine-Google; (+http://code.google.com/appengine)'
        self.request_ip = '0.1.0.1'
        self.request_host_name = 'py3.aou-res-curation-test.appspot.com'
        self.request_log_id = 'fake_request_id'
        self.request_trace_id = 'fake_trace_id'
        self.request_trace = 'projects/{0}/traces/{1}'.format(
            self.project_id, self.request_trace_id)

        self.request_start_time = pytz.utc.localize(datetime(2020, 1, 1))
        self.request_end_time = pytz.utc.localize(datetime(
            2020, 1, 1)) + timedelta(minutes=1)
        self.request_log_entry_ts = self.request_start_time + timedelta(
            seconds=10)
        self.log_record_created = self.request_start_time - timedelta(
            seconds=10)

        self.mock_get_application_id_patcher = patch(
            'app_identity.get_application_id')
        self.mock_get_application_id = self.mock_get_application_id_patcher.start(
        )
        self.mock_get_application_id.return_value = self.project_id

        # Mock a flask request for testing
        self.request = MagicMock()
        type(self.request).method = PropertyMock(
            return_value=self.request_method)
        type(self.request).full_path = PropertyMock(
            return_value=self.request_full_path)
        type(self.request).user_agent = PropertyMock(
            return_value=self.request_user_agent)

        headers = dict()
        headers['X-Appengine-User-Ip'] = self.request_ip
        headers['X-Appengine-Default-Version-Hostname'] = self.request_host_name
        headers['X-Appengine-Request-Log-Id'] = self.request_log_id
        headers['X-Appengine-Taskname'] = None
        headers['X-Appengine-Queuename'] = None
        headers['X-Cloud-Trace-Context'] = self.request_trace_id

        type(self.request).headers = PropertyMock(return_value=headers)

        # Define the log records for testing
        self.file_path = 'data_steward/validation/main'
        self.file_name = 'main'
        self.info_log_record = self.create_log_record(
            'info', self.log_record_created, logging.INFO, self.file_name,
            self.file_path, 10, 'info message')
        self.debug_log_record = self.create_log_record(
            'debug', self.log_record_created, logging.DEBUG, self.file_name,
            self.file_path, 11, 'debug message')
        self.error_log_record = self.create_log_record(
            'error', self.log_record_created, logging.ERROR, self.file_name,
            self.file_path, 12, 'error message')

        self.info_log_line = {
            'logMessage': 'info message',
            'severity': SEVERITY_INFO,
            'time': self.log_record_created.isoformat(),
            'sourceLocation': {
                'file': self.file_path,
                'functionName': self.file_name,
                'line': 10
            }
        }

        self.debug_log_line = {
            'logMessage': 'debug message',
            'severity': SEVERITY_DEBUG,
            'time': self.log_record_created.isoformat(),
            'sourceLocation': {
                'file': self.file_path,
                'functionName': self.file_name,
                'line': 11
            }
        }

        self.error_log_line = {
            'logMessage': 'error message',
            'severity': SEVERITY_ERROR,
            'time': self.log_record_created.isoformat(),
            'sourceLocation': {
                'file': self.file_path,
                'functionName': self.file_name,
                'line': 12
            }
        }

        self.mock_logging_service_client_patcher = patch(
            'curation_logging.curation_gae_handler.gcp_logging_v2.LoggingServiceV2Client'
        )
        self.mock_logging_service_client = self.mock_logging_service_client_patcher.start(
        )

    def tearDown(self):
        self.mock_logging_service_client_patcher.stop()
        self.request.stop()

    @staticmethod
    def create_log_record(name, record_created, level_no, func_name, pathname,
                          lineno, msg):
        log_record = LogRecord(name=name,
                               levelno=level_no,
                               lineno=lineno,
                               func=func_name,
                               pathname=pathname,
                               msg=msg,
                               level=level_no,
                               args={},
                               exc_info={})
        log_record.created = record_created
        return log_record

    @mock.patch('curation_logging.curation_gae_handler.datetime')
    def test_gcp_stackdriver_logger(self, mock_datetime):
        mock_datetime.now.return_value.isoformat.return_value = self.request_start_time.isoformat(
        )
        mock_datetime.utcnow.return_value = self.request_start_time
        mock_datetime.utcfromtimestamp.return_value = self.log_record_created

        # Initialize GCPStackDriverLogger
        self.gcp_stackdriver_logger = GCPStackDriverLogger(LOG_BUFFER_SIZE)
        self.gcp_stackdriver_logger.setup_from_request(self.request)

        self.assertIsNone(self.gcp_stackdriver_logger._first_log_ts)
        self.assertEqual(self.gcp_stackdriver_logger._start_time,
                         self.request_start_time.isoformat())
        self.assertEqual(self.gcp_stackdriver_logger._request_method,
                         self.request_method)
        self.assertEqual(self.gcp_stackdriver_logger._request_resource,
                         self.request_full_path)
        self.assertEqual(self.gcp_stackdriver_logger._request_agent,
                         self.request_user_agent)
        self.assertEqual(self.gcp_stackdriver_logger._request_remote_addr,
                         self.request_ip)
        self.assertEqual(self.gcp_stackdriver_logger._request_host,
                         self.request_host_name)
        self.assertEqual(self.gcp_stackdriver_logger._request_log_id,
                         self.request_log_id)
        self.assertEqual(self.gcp_stackdriver_logger._trace, self.request_trace)

        self.gcp_stackdriver_logger.log_event(self.info_log_record)
        self.gcp_stackdriver_logger.log_event(self.debug_log_record)
        self.gcp_stackdriver_logger.log_event(self.error_log_record)

        self.assertEqual(
            len(self.gcp_stackdriver_logger._buffer), 0,
            'expected log buffer to flush itself after being filled')
        self.assertEqual(
            self.mock_logging_service_client.return_value.write_log_entries.
            call_count, 1)

        self.gcp_stackdriver_logger.finalize()
        self.assertIsNone(self.gcp_stackdriver_logger._first_log_ts)
        self.assertEqual(self.gcp_stackdriver_logger._start_time, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_method, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_resource, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_agent, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_remote_addr, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_host, None)
        self.assertEqual(self.gcp_stackdriver_logger._request_log_id, None)
        self.assertEqual(self.gcp_stackdriver_logger._trace, None)

    @mock.patch('curation_logging.curation_gae_handler.get_gcp_logger')
    def test_initialize_logging(self, mock_get_gcp_logger):
        with patch.dict('os.environ', {'GAE_ENV': ''}):
            curation_gae_handler.initialize_logging(logging.DEBUG)
            logging.info(self.info_log_record)
            logging.debug(self.debug_log_line)
            logging.error(self.error_log_record)
            self.assertEqual(3, mock_get_gcp_logger.call_count)

    @mock.patch('requests.get')
    def test_setup_logging_zone(self, mock_requests_get):
        with patch.dict('os.environ', {'GAE_SERVICE': ''}):
            timezone = 'test time zone'
            mock_response = MagicMock()
            type(mock_response).status_code = PropertyMock(return_value=200)
            type(mock_response).text = timezone
            mock_requests_get.return_value = mock_response
            actual_time_zone = curation_gae_handler.setup_logging_zone()
            self.assertEqual(actual_time_zone, timezone)

            type(mock_response).status_code = PropertyMock(return_value=500)
            actual_time_zone = curation_gae_handler.setup_logging_zone()
            self.assertEqual(actual_time_zone, 'local-machine')

        with patch.dict('os.environ', dict()):
            actual_time_zone = curation_gae_handler.setup_logging_zone()
            self.assertEqual(actual_time_zone, 'local-machine')

    @mock.patch('curation_logging.curation_gae_handler.setup_logging_zone')
    def test_setup_logging_resource(self, mock_setup_logging_zone):
        timezone = 'test time zone'
        mock_setup_logging_zone.return_value = timezone
        actual_resource = curation_gae_handler.setup_logging_resource()
        expected_resource = MonitoredResource(
            type='gae_app',
            labels={
                'project_id': self.project_id,
                'module_id': GAE_LOGGING_MODULE_ID,
                'version_id': GAE_LOGGING_VERSION_ID,
                'zone': timezone
            })
        self.assertEqual(expected_resource, actual_resource)

    @mock.patch('curation_logging.curation_gae_handler.datetime')
    @mock.patch(
        'curation_logging.curation_gae_handler.gcp_logging._helpers._normalize_severity'
    )
    def test_setup_log_line(self, mock_normalize_severity, mock_datetime):
        mock_datetime.utcfromtimestamp.return_value = self.log_record_created
        mock_normalize_severity.side_effect = [
            SEVERITY_INFO, SEVERITY_DEBUG, SEVERITY_ERROR
        ]

        actual_info_log_line = curation_gae_handler.setup_log_line(
            self.info_log_record)
        self.assertDictEqual(self.info_log_line, actual_info_log_line)

        actual_debug_log_line = curation_gae_handler.setup_log_line(
            self.debug_log_record)
        self.assertDictEqual(self.debug_log_line, actual_debug_log_line)

        actual_error_log_line = curation_gae_handler.setup_log_line(
            self.error_log_record)
        self.assertDictEqual(self.error_log_line, actual_error_log_line)

    def test_get_highest_severity_level_from_lines(self):
        lines = [self.info_log_line, self.debug_log_line, self.error_log_line]
        actual_severity_level = curation_gae_handler.get_highest_severity_level_from_lines(
            lines)
        self.assertEqual(SEVERITY_ERROR, actual_severity_level)

    def test_setup_proto_payload(self):
        lines = [self.info_log_line, self.debug_log_line, self.error_log_line]

        proto_payload_args = {
            'startTime': self.request_start_time.isoformat(),
            'endTime': self.request_end_time.isoformat(),
            'method': self.request_method,
            'resource': self.request_full_path,
            'userAgent': self.request_user_agent,
            'host': self.request_host_name,
            'ip': self.request_ip,
            'responseSize': None
        }

        actual_proto_payload = curation_gae_handler.setup_proto_payload(
            lines, LogCompletionStatusEnum.PARTIAL_BEGIN, **proto_payload_args)

        expected_dict = dict(
            {
                '@type': curation_gae_handler.REQUEST_LOG_TYPE,
                'first': True,
                'finished': False,
                'line': lines
            }, **proto_payload_args)

        expected_proto_payload = gcp_json_format.ParseDict(
            expected_dict, gcp_any_pb2.Any())

        self.assertEqual(expected_proto_payload, actual_proto_payload)

    def test_update_long_operation(self):
        expected_operation = LogEntryOperation(
            id=self.request_log_id,
            producer='appengine.googleapis.com/request_id',
            first=True,
            last=True)

        actual_operation = curation_gae_handler.update_long_operation(
            self.request_log_id, LogCompletionStatusEnum.COMPLETE)
        self.assertEqual(expected_operation, actual_operation)

        expected_operation = LogEntryOperation(
            id=self.request_log_id,
            producer='appengine.googleapis.com/request_id',
            first=True,
            last=False)

        actual_operation = curation_gae_handler.update_long_operation(
            self.request_log_id, LogCompletionStatusEnum.PARTIAL_BEGIN)
        self.assertEqual(expected_operation, actual_operation)

        expected_operation = LogEntryOperation(
            id=self.request_log_id,
            producer='appengine.googleapis.com/request_id',
            first=False,
            last=False)

        actual_operation = curation_gae_handler.update_long_operation(
            self.request_log_id, LogCompletionStatusEnum.PARTIAL_MORE)
        self.assertEqual(expected_operation, actual_operation)


# -*- coding: utf-8 -*-
# Felipe G. Foschiera
# URI 2018 - Olimpíadas de Natal
from collections import defaultdict
medalhas = defaultdict(list)
while True:
    try:
        categoria, ouro, prata, bronze = input(), input(), input(), input()
        if ouro not in medalhas:
            medalhas[ouro] = [1, 0, 0]
        else:
            medalhas[ouro][0] += 1
        if prata not in medalhas:
            medalhas[prata] = [0, 1, 0]
        else:
            medalhas[prata][1] += 1
        if bronze not in medalhas:
            medalhas[bronze] = [0, 0, 1]
        else:
            medalhas[bronze][2] += 1
    except EOFError:
        break
print("Quadro de Medalhas")
for m in sorted(medalhas.items(), key = lambda x: (-x[1][0], -x[1][1], -x[1][2], x[0])):
    print(m[0], m[1][0], m[1][1], m[1][2])

URI Online Judge/Estruturas e Bibliotecas/2018 - Olimpíadas de Natal/natal.py


from liir.ml.core.layers.Layer import Layer
import theano as th
__author__ = 'quynhdo'

class Reverse(Layer):
    # reverse the input
    def __init__(self, idx="0",  times=2):
        Layer.__init__(self, id=idx)
        self.times = times

    def compile(self):
        axes = [self.input.ndim - 1]
        slices = [slice(None, None, -1) if i in axes else slice(None, None, None) for i in range(self.input.ndim)]
        self.output = self.input[slices]


import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import  MIMEText
from email.mime.base import MIMEBase
from email import encoders
import os

def sendMail(columnames,prop):
    fromaddr = "contacto@bullestate.cl"
    toaddr = "joaquin@bullestate.cl, sergei@bullestate.cl"

    msg = MIMEMultipart()

    msg['From'] = fromaddr
    msg['To'] = toaddr
    msg['Subject'] = "Oportunidad BullEstate"

    body = "Estimados" + ":\n\nOportunidad Encontrada."
    for i,x in enumerate(columnames):
        body+=(str(columnames[i]))
        body+=": "
        try:
            body+=(str(prop[i]))
        except:
            body+=""
        body+='\n'

    msg.attach(MIMEText(body, 'plain'))

    server = smtplib.SMTP_SSL('smtp.zoho.com', 465)
    #server.starttls()
    server.login(fromaddr, "Bullestate.123")
    text = msg.as_string()
    server.sendmail(fromaddr, toaddr, text)
    server.quit()