# Test performance evaluator

from definitions import *
from pyopus.evaluator.auxfunc import listParamDesc, paramDict
from pyopus.evaluator.performance import PerformanceEvaluator
# If MPI is imported an application not using MPI will behave correctly
# (i.e. only slot 0 will run the program) even when started with mpirun/mpiexec
from pyopus.parallel.mpi import MPI
from pyopus.parallel.cooperative import cOS
# Plotting support
from pyopus.plotter import interface as pyopl
# pprint for pretty printing data structures
from pprint import pprint
# pickle for loading results
import pickle, os
# NumPy
import numpy as np


if __name__=='__main__':
	# Prepare statistical parameters dictionary with nominal values (0)
	nominalStat=paramDict(np.zeros(len(statParams)), list(statParams.keys()))
	
	# Prepare operating parameters dictionary with nominal values
	names=list(opParams.keys())
	nominalOp=paramDict(listParamDesc(opParams, names, "init"), names)
	
	# Prepare nominal design parameters dictionaries
	names=list(designParams.keys())
	nominalDesign=paramDict(listParamDesc(designParams, names, "init"), names)
	
	# Prepare one corner, module 'tm', nominal op parameters
	corners={
		'nom': {
			'heads': [ 'opus' ], 
			'modules': ['tm'], 
			'params': nominalOp, 
			
		}
	}
	
	# Prepare parallel environment
	cOS.setVM(MPI(mirrorMap={'*':'.'}))
	
	# Measures have no corners listed - they are evaluated across all specified corners
	pe=PerformanceEvaluator(
		heads, analyses, measures, corners, variables=variables, 
		storeResults=True, resultsFolder=None, resultsPrefix="restemp.", 
		debug=2
	)
	results, anCount = pe([nominalDesign, nominalStat])
	print(pe.formatResults(nMeasureName=10, nCornerName=15))
	
	# Print the list of temporary pickled result files across hosts where 
	# evaluations took place. This is a dictionary with (host, (corner, analysis))
	# for key. None for host means the results are stored on the local machine
	# (i.e. the one that invoked the evaluator). 
	print("\nResult files across hosts")
	pprint(pe.resFiles)
	
	# Extract and print the set of hosts where evaluations took place
	s=set()
	for h, k in pe.resFiles.keys():
		s.add(h)
	print("\nSet of hosts where results are stored (None is localhost):", s)
	
	# To access all result files we must first move them to the host where 
	# pe() was called. The files will be moved to the current directory 
	# (destination='.'). The collected files will have prefix 'res_'. Files 
	# will be moved (i.e. deleted from remote machines). 
	files={}
	files=pe.collectResultFiles(destination=".", prefix="res.", move=True)
	
	# Print names of pickled result files after they were moved here. 
	# This is a dictionary with (corner, analysis) for key. 
	print("\nCollected result files")
	pprint(files)
	
	# Result file with key ('nom', 'blank') in the above mentioned dictionary 
	# holds the environment where dependent measurements without were 
	# evaluated for the 'nom' corner. Load one and print the environment. 
	with open(files[('nom', 'blank')], 'rb') as f:
		noneRes=pickle.load(f)
	
	print("\n'blank' analysis measure evaluation environment")
	pprint(noneRes.evalEnvironment())
	
	# Load the ac analysis results and plot gain in dB vs frequency
	with open(files[('nom', 'ac')], 'rb') as f:
		acRes=pickle.load(f)
	
	# Extract frequency scale and gain. 
	freq=np.real(acRes.scale())
	tf=acRes.v('out')/acRes.v('inp', 'inn')
	gain=20*np.log10(np.abs(tf))
	ph=np.unwrap(np.angle(tf))*180/np.pi

	# Plot gain. 
	f1=pyopl.figure(windowTitle="AC analysis", figpx=(600,400), dpi=100)	
	pyopl.lock(True)
	if pyopl.alive(f1):
		ax1=f1.add_axes((0.12,0.12,0.76,0.76))
		ax1.semilogx(freq, gain, '-', label='gain [dB]', color=(1,0,0))
		ax1.semilogx(freq, ph, '-', label='phase [deg]', color=(0,0.5,0))
		ax1.set_xlabel('f [Hz]')
		ax1.set_ylabel('gain [dB]')
		ax1.set_title('Gain vs frequency')
		ax1.legend(loc='upper right')
		ax1.grid(True)
		pyopl.draw(f1)
	pyopl.lock(False)
	
	# Wait for the user to close the plot environment control window. 
	pyopl.join()

	# Delete collected result files
	for k, f in files.items():
		os.remove(f)
	
	# Cleanup temporary files generated by the evaluator
	pe.finalize()
	
	# Finalize cOS parallel environment
	cOS.finalize()