import os
from MatlabFuncs import *
from expandvariables import *
from expandresponses import *
from helpers import *
from dakota_in_data import *
from process_qmu_response_data import *


def preqmu(md, options):
    '''QMU - apply Quantification of Margins and Uncertainties techniques
    to a solution sequence (like stressbalance.py, progonstic.py, etc ...),
    using the Dakota software from Sandia.

   options come from the solve.py routine. They can include Dakota options:

    qmufile: input file for Dakota

    (ivap, iresp, imethod, and iparams are currently unimplemented)
    ivar: selection number for variables input (if several are specified in variables)
    iresp: same thing for response functions
    imethod: same thing for methods
    iparams: same thing for params
'''

    print('preprocessing dakota inputs')
    qmufile = options.getfieldvalue('qmufile', 'qmu')
    # qmufile cannot be changed unless ????script.sh is also changed
    options.addfielddefault('ivar', 0)
    options.addfielddefault('iresp', 0)
    options.addfielddefault('imethod', 0)
    options.addfielddefault('iparams', 0)

    # when running in library mode, the in file needs to be called md.miscellaneous.name.qmu.in
    qmufile = md.miscellaneous.name

    # retrieve variables and resposnes for this particular analysis.
    #print type(md.qmu.variables)
    #print md.qmu.variables.__dict__
    #print ivar
    variables = md.qmu.variables  #[ivar]
    responses = md.qmu.responses  #[iresp]

    # expand variables and responses
    #print variables.__dict__
    #print responses.__dict__
    variables = expandvariables(md, variables)
    responses = expandresponses(md, responses)

    # go through variables and responses, and check they don't have more than
    #   md.qmu.numberofpartitions values. Also determine numvariables and numresponses
    #[[[
    numvariables = 0
    variable_fieldnames = fieldnames(variables)
    for i in range(len(variable_fieldnames)):
        field_name = variable_fieldnames[i]
        fieldvariables = vars(variables)[field_name]
        for j in range(np.size(fieldvariables)):
            if strncmpi(fieldvariables[j].descriptor, '\'scaled_', 8) and str2int(fieldvariables[j].descriptor, 'last') > md.qmu.numberofpartitions:
                raise RuntimeError('preqmu error message: one of the expanded variables has more values than the number of partitions (setup in md.qmu.numberofpartitions)')

        numvariables = numvariables + np.size(vars(variables)[field_name])

    numresponses = 0
    response_fieldnames = fieldnames(responses)
    for i in range(len(response_fieldnames)):
        field_name = response_fieldnames[i]
        fieldresponses = vars(responses)[field_name]
        for j in range(np.size(fieldresponses)):
            if strncmpi(fieldresponses[j].descriptor, '\'scaled_', 8) and str2int(fieldresponses[j].descriptor, 'last') > md.qmu.numberofpartitions:
                raise RuntimeError('preqmu error message: one of the expanded responses has more values than the number of partitions (setup in md.qmu.numberofpartitions)')

        numresponses = numresponses + np.size(vars(responses)[field_name])

    #]]]

    # create in file for dakota
    #dakota_in_data(md.qmu.method[imethod], variables, responses, md.qmu.params[iparams], qmufile)
    dakota_in_data(md.qmu.method, variables, responses, md.qmu.params, qmufile)

    # build a list of variables and responses descriptors. the list is not expanded.
    #{{{
    variabledescriptors = []
    # variable_fieldnames = fieldnames(md.qmu.variables[ivar])
    variable_fieldnames = fieldnames(md.qmu.variables)
    for i in range(len(variable_fieldnames)):
        field_name = variable_fieldnames[i]
    #fieldvariables = vars(md.qmu.variables[ivar])[field_name]
        fieldvariables = vars(md.qmu.variables)[field_name]
        if type(fieldvariables) in [list, np.ndarray]:
            for j in range(np.size(fieldvariables)):
                variabledescriptors.append(fieldvariables[j].descriptor)
        else:
            variabledescriptors.append(fieldvariables.descriptor)

    responsedescriptors = []
    # response_fieldnames = fieldnames(md.qmu.responses[iresp])
    response_fieldnames = fieldnames(md.qmu.responses)
    for i in range(len(response_fieldnames)):
        field_name = response_fieldnames[i]
    #fieldresponses = vars(md.qmu.responses[iresp])[field_name]
        fieldresponses = vars(md.qmu.responses)[field_name]
        for j in range(np.size(fieldresponses)):
            responsedescriptors.append(fieldresponses[j].descriptor)
    #}}}

    # register the fields that will be needed by the Qmu model.
    md.qmu.numberofresponses = numresponses
    md.qmu.variabledescriptors = variabledescriptors
    md.qmu.responsedescriptors = responsedescriptors

    # now, we have to provide all the info necessary for the solutions to compute the
    # responses. For ex, if mass_flux is a response, we need a profile of points.
    # For a misfit, we need the observed velocity, etc ...
    md = process_qmu_response_data(md)

    return md
