import numpy as np
from MatlabFuncs import *
from IssmConfig import *
from project3d import project3d
from collections import OrderedDict
from fielddisplay import fielddisplay
from checkfield import checkfield
from WriteData import WriteData
from helpers import *
from dakota_method import *

class qmu(object):
	"""
	QMU class definition

	   Usage:
	      qmu=qmu();
	"""

	def __init__(self): # {{{
		self.isdakota                    = 0
		self.variables                   = OrderedStruct()
		self.responses                   = OrderedStruct()
		self.method                      = OrderedDict()
		self.params                      = OrderedStruct()
		self.results                     = OrderedDict()
		self.partition                   = float('NaN')
		self.numberofpartitions          = 0
		self.numberofresponses           = 0
		self.variabledescriptors         = []
		self.responsedescriptors         = []
		self.mass_flux_profile_directory = float('NaN')
		self.mass_flux_profiles          = float('NaN')
		self.mass_flux_segments          = []
		self.adjacency                   = float('NaN')
		self.vertex_weight               = float('NaN')

		#set defaults
		self.setdefaultparameters()

		#}}}
	def __repr__(self):    # {{{
		s ='   qmu parameters:\n'

		s+="%s\n" % fielddisplay(self,'isdakota','is qmu analysis activated?')
		maxlen = 0
		s+="         variables:  (arrays of each variable class)\n"

		# OrderedStruct's iterator returns individual name/array-of-functions pairs
		for variable in self.variables:
			fname=variable[0]
			maxlen=max(maxlen,len(fname))
			size = np.shape(variable[1])
			a = size[0]
			b = 1 if len(size) < 2 else size[1]
			s+="            %-*s:    [%ix%i]    '%s'\n" %  (maxlen+1,fname,a,b,type(variable[1][0]))

		s+="         responses:  (arrays of each response class)\n"
		for response in self.responses:
			fname=response[0]
			maxlen=max(maxlen,len(fname))
			size = np.shape(response[1])
			a = size[0]
			b = 1 if len(size) < 2 else size[1]
			s+="            %-*s:    [%ix%i]    '%s'\n" %  (maxlen+1,fname,a,b,type(response[1][0]))

		s+="%s\n" % fielddisplay(self,'numberofresponses','number of responses')

		if type(self.method) != OrderedDict:
			self.method = [self.method]
		# self.method must be iterable
		for method in self.method:
			if isinstance(method,dakota_method):
				s+="            method :    '%s'\n" % (method.method)

		# params could be have a number of forms (mainly 1 struct or many)
		if type(self.params) == OrderedStruct:
			params = [self.params]
		else:
			params = np.hstack(np.atleast_1d(np.array(self.params)))
		for param in params:
			print(type(param))
			print(param)
			s+="         params:  (array of method-independent parameters)\n"
			fnames=vars(param)
			maxlen=0
			for fname in fnames:
				maxlen=max(maxlen,len(fname))

			for fname in fnames:
				s+="            %-*s: %s\n" %  (maxlen+1,fname,str(getattr(param,fname)))

		# results could be have a number of forms (mainly 1 struct or many)
		results = np.hstack(np.atleast_1d(np.array(self.results)))
		for result in results:
			s+="         results:  (information from dakota files)\n"
			fnames=vars(result)
			maxlen=0
			for fname in fnames:
				maxlen=max(maxlen,len(fname))

			for fname in fnames:
				size = np.shape(response[1])
				a = size[0]
				b = 0 if len(size) < 2 else size[1]
				size = np.shape(getattr(result,fname))
				s+="            %-*s:    [%ix%i]    '%s'\n" % (maxlen+1,fname,a,b,type(getattr(result,fname)))

		s+="%s\n" % fielddisplay(self,'partition','user provided mesh partitioning, defaults to metis if not specified') 
		s+="%s\n" % fielddisplay(self,'numberofpartitions','number of partitions for semi-discrete qmu') 
		s+="%s\n" % fielddisplay(self,'variabledescriptors','')
		s+="%s\n" % fielddisplay(self,'responsedescriptors','')
		s+="%s\n" % fielddisplay(self,'method','array of dakota_method class')
		s+="%s\n" % fielddisplay(self,'mass_flux_profile_directory','directory for mass flux profiles')
		s+="%s\n" % fielddisplay(self,'mass_flux_profiles','list of mass_flux profiles')
		s+="%s\n" % fielddisplay(self,'mass_flux_segments','')
		s+="%s\n" % fielddisplay(self,'adjacency','')
		s+="%s\n" % fielddisplay(self,'vertex_weight','weight applied to each mesh vertex')

		return s
	# }}}
	def extrude(self,md): # {{{
		self.partition=project3d(md,'vector',np.transpose(self.partition),'type','node')
		return self
	#}}}
	def setdefaultparameters(self): # {{{
		return self
	#}}}
	def checkconsistency(self,md,solution,analyses):    # {{{

		#Early return
		if not md.qmu.isdakota:
			return

		version=IssmConfig('_DAKOTA_VERSION_')
		version=float(version[0])

		if version < 6:
			if not md.qmu.params.evaluation_concurrency==1:
				md.checkmessage("concurrency should be set to 1 when running dakota in library mode")
		else:
			if not strcmpi(self.params.evaluation_scheduling,'master'):
				md.checkmessage('evaluation_scheduling in qmu.params should be set to "master"')

			if md.cluster.np <= 1:
				md.checkmessage('in parallel library mode, Dakota needs to run on at least 2 cpus, 1 cpu for the master, 1 cpu for the slave. Modify md.cluser.np accordingly.')
					
			if self.params.processors_per_evaluation < 1:
				md.checkmessage('in parallel library mode, Dakota needs to run at least one slave on one cpu (md.qmu.params.processors_per_evaluation >=1)!')
				
			if np.mod(md.cluster.np-1,self.params.processors_per_evaluation):
				md.checkmessage('in parallel library mode, the requirement is for md.cluster.np = md.qmu.params.processors_per_evaluation * number_of_slaves, where number_of_slaves will automatically be determined by Dakota. Modify md.cluster.np accordingly')
		
		if np.size(md.qmu.partition) > 0:
			if np.size(md.qmu.partition)!=md.mesh.numberofvertices and np.size(md.qmu.partition) != md.mesh.numberofelements:
				md.checkmessage("user supplied partition for qmu analysis should have size (md.mesh.numberofvertices x 1) or (md.mesh.numberofelements x 1)")
			if not min(md.qmu.partition.flatten())==0:
				md.checkmessage("partition vector not indexed from 0 on")
			if max(md.qmu.partition.flatten())>=md.qmu.numberofpartitions:
				md.checkmessage("for qmu analysis, partitioning vector cannot go over npart, number of partition areas")

		return md
	# }}}
	def marshall(self,prefix,md,fid):    # {{{
		WriteData(fid,prefix,'object',self,'fieldname','isdakota','format','Boolean')
		if not self.isdakota:
			WriteData(fid,prefix,'data',False,'name','md.qmu.mass_flux_segments_present','format','Boolean');
			return
		WriteData(fid,prefix,'object',self,'fieldname','partition','format','DoubleMat','mattype',2)
		WriteData(fid,prefix,'object',self,'fieldname','numberofpartitions','format','Integer')
		WriteData(fid,prefix,'object',self,'fieldname','numberofresponses','format','Integer')
		WriteData(fid,prefix,'object',self,'fieldname','variabledescriptors','format','StringArray')
		WriteData(fid,prefix,'object',self,'fieldname','responsedescriptors','format','StringArray')
		if not isempty(self.mass_flux_segments):
			WriteData(fid,prefix,'data',self.mass_flux_segments,'name','md.qmu.mass_flux_segments','format','MatArray');
			flag=True; 
		else:
			flag=False; 
		WriteData(fid,prefix,'data',flag,'name','md.qmu.mass_flux_segments_present','format','Boolean');
	# }}}