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.output   = 0
        self.variables = OrderedStruct()
        self.responses = OrderedStruct()
        self.method = OrderedDict()
        self.params = OrderedStruct()
        self.results = OrderedDict()
        self.vpartition = float('NaN')
        self.epartition = 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?')
        s += "%s\n" % fielddisplay(self, 'output', 'are we outputting ISSM results, default is 0')
        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, 'vpartition', 'user provided mesh partitioning (vertex based)')
        s += "%s\n" % fielddisplay(self, 'epartition', 'user provided mesh partitioning (element based)')
        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.vpartition = project3d(md, 'vector', np.transpose(self.vpartition), 'type', 'node')
        self.epartition = project3d(md, 'vector', np.transpose(self.epartition), 'type', 'element')
        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.cluster.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.vpartition) > 0:
            if np.size(md.qmu.vpartition, 0) != md.mesh.numberofvertices:
                md.checkmessage("user supplied vertex partition for qmu analysis should have size (md.mesh.numberofvertices x 1)")
            if not min(md.qmu.vpartition.flatten()) == 0:
                md.checkmessage("vertex partition vector not indexed from 0 on")
            if max(md.qmu.vpartition.flatten()) >= md.qmu.numberofpartitions:
                md.checkmessage("for qmu analysis, vertex partitioning vector cannot go over npart, number of partition areas")

                if np.size(md.qmu.epartition) > 0:
                    if np.size(md.qmu.epartition, 0) != md.mesh.numberofelements:
                        md.checkmessage("user supplied element partition for qmu analysis should have size (md.mesh.numberofelements x 1)")
                    if not min(md.qmu.epartition.flatten()) == 0:
                        md.checkmessage("elememtn partition vector not indexed from 0 on")
                    if max(md.qmu.epartition.flatten()) >= md.qmu.numberofpartitions:
                        md.checkmessage("for qmu analysis, element partitioning vector cannot go over npart, number of partition areas")

                if np.size(md.qmu.vpartition) == 0 or np.any(np.isnan(md.qmu.vpartition)) or np.size(md.qmu.epartition) == 0 or np.any(np.isnan(md.qmu.epartition)):
                    md.checkmessage("for qmu analysis, both an element and partitioning vectors need to be supplied with no nan values! One can be defaulted to all zeros.")

        return md
    # }}}
    def marshall(self, prefix, md, fid):  # {{{
        WriteData(fid, prefix, 'object', self, 'fieldname', 'isdakota', 'format', 'Boolean')
        WriteData(fid, prefix, 'object', self, 'fieldname', 'output', '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', 'vpartition', 'format', 'DoubleMat', 'mattype', 2)
        WriteData(fid, prefix, 'object', self, 'fieldname', 'epartition', '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')
    # }}}
