#Move this somewhere else later
from helpers import *

from dakota_method import *
from dakota_in_params import *
from IssmConfig import *
from dmeth_params_write import *
from vector_write import *
from qmu_classes import *

import itertools


def dakota_in_write(method, dvar, dresp, params, filei, *args):
    """DAKOTA_IN_WRITE - Write a Dakota .in input file.

    Usage:
        [] = dakota_in_write(method, dvar, dresp, params, filei, args)
        [] = dakota_in_write(dmeth , dvar, dresp, params, filei, args)

    where the required input is,
        method        (character, Dakota method name)
        dmeth         (dakota_method, method class object)
        dvar          (structure array, variable class objects)
        dresp         (structure array, response class objects)
        params        (structure array, method-independent parameters)
        filei         (character, name of .in file)

    The method and filei will be prompted for if empty. params may be empty, in 
    which case defaults will be used.

    The optional args are not yet used.

    This function writes a Dakota .in input file to be used by Dakota. This 
    file is indepent of the particular analysis package.

    This data would typically be generated by a matlab script for a specific 
    model, using the method, variable, and response class objects. This 
    function may be called by dakota_in_data.

    "Copyright 2009, by the California Institute of Technology. ALL RIGHTS 
    RESERVED. United States Government Sponsorship acknowledged. Any commercial 
    use must be negotiated with the Office of Technology Transfer at the 
    California Institute of Technology. (NTR 47078)

    This software may be subject to U.S. export control laws. By accepting this 
    software, the user agrees to comply with all applicable U.S. export laws 
    and regulations. User has the responsibility to obtain export licenses, or 
    other export authority as may be required before exporting such information 
    to foreign countries or providing access to foreign persons."
    """

    #  process the input parameters
    if len(fieldnames(method)) == 0:
        method = str(eval(input('Method?  ')))

    if type(method) == str:
        dmeth = dakota_method(method)
    elif isinstance(method, dakota_method):
        dmeth = method
    else:
        raise RuntimeError('Method ' + str(method) + ' is unrecognized class ' + str(type(method)) + '. (should be either "str" or "dakota_method")')

    if len(filei) == 0:
        filei = str(eval(input('Dakota input file to write?  ')))

    pathstr, name, ext = fileparts(filei)
    if len(ext) == 0:
        # fileparts only considers '.in' to be the extension, not '.qmu.in'
        ext = '.qmu.in'

    filei2 = fullfile(pathstr, name + ext)

    print('Opening Dakota input file \'' + filei2 + '\'')
    try:
        with open(filei2, 'w+') as fidi:

            if len(fieldnames(params)) == 0:
                params = struct()

            params = dakota_in_params(params)

            #  write the strategy section
            if float(IssmConfig('_DAKOTA_VERSION_')[0]) < 6:
                strategy_write(fidi, params)
            else:
                environment_write(fidi, params)

            #  write the method section
            method_write(fidi, dmeth, dresp, params)
            #  write the model section
            model_write(fidi)
            #  write the variables section
            variables_write(fidi, dmeth, dvar)
            #  write the interface section
            interface_write(fidi, params)
            #  write the responses section
            responses_write(fidi, dmeth, dresp, params)

    except IOError:
        print(filei2 + ' could not be opened')

    print('End of file successfully written')

    # Uncomment to print contents of Dakota input file (for debugging)
    # with open(filei2, 'r') as fidi:
    #     print(fidi.read())

#  function to write the strategy section of the file
def strategy_write(fidi, params):

    print('Writing strategy section of Dakota input file')

    fidi.write('strategy, \n')
    fidi.write('\tsingle_method\n\n')
    param_write(fidi, '\t  ', 'graphics', '', '\n', params)
    param_write(fidi, '\t  ', 'tabular_graphics_data', '', '\n', params)
    param_write(fidi, '\t  ', 'tabular_graphics_file', ' ', '\n', params)
    fidi.write('\n')


#  function to write the environment section of the file
def environment_write(fidi, params):

    print('Writing environment section of Dakota input file')

    fidi.write('environment, \n')
    param_write(fidi, '\t  ', 'graphics', '', '\n', params)
    param_write(fidi, '\t  ', 'tabular_graphics_data', '', '\n', params)
    param_write(fidi, '\t  ', 'tabular_graphics_file', ' ', '\n', params)
    fidi.write('\n')


#  function to write the method section of the file
def method_write(fidi, dmeth, dresp, params):

    print('Writing method section of Dakota input file')

    fidi.write('method, \n')
    fidi.write('\t' + str(dmeth.method) + '\n')

    dmeth_params_write(dmeth, fidi)

    #  write response levels

    if strcmp(dmeth.type, 'nond'):
        for i in range(len(dmeth.responses)):
            str_name = dmeth.responses[i]
            resp = eval("{}.{}()".format(str_name, str_name))
            resp.dakota_rlev_write(fidi, dresp, params)

    fidi.write('\n')


#  function to write the model section of the file
def model_write(fidi):

    print('Writing model section of Dakota input file')

    fidi.write('model, \n')
    fidi.write('\tsingle\n\n')


#  function to write the variables section of the file
def variables_write(fidi, dmeth, dvar):

    # print('Writing variables section of Dakota input file')

    # fidi.write('variables, \n')

    # #  variables vary by method
    # fd = fieldnames(dvar)
    # types = []
    # var = []
    # for i in range(len(fd)):
    #     i_type = eval('dvar.{}[0].__class__.__name__'.format(fd[i]))
    #     j = dmeth.variables.index(i_type)
    #     str_name = dmeth.variables[j]

    # # organize so that multiple instances of the same qmu class
    # # (2 different variable instances of "normal_uncertain" for example)
    # # are in the same dakota_write call regardless of individual size
    # # but that each class has its own dakota_write call
    #     if str_name not in types:
    #         types.append(str_name)
    #         var.append(eval('dvar.{}'.format(fd[i])))
    #     else:
    #         t = types.index(str_name)
    #         var[t].extend(eval('dvar.{}'.format(fd[i])))

    # for t in range(len(types)):
    #     v = eval('{}.{}()'.format(types[t], types[t]))
    #     v.dakota_write(fidi, var[t])

    # #  linear constraints vary by method
    # fc = dmeth.lcspec

    # for i in range(len(dmeth.lcspec)):
    #     str_name = dmeth.lcspec[i]
    #     var = eval('{}.{}()'.format(str_name, str_name))
    # # check that str_name is correct against matlab version which has no argument there
    #     var.dakota_write(fidi, eval('dvar.{}[i]'.format(j)), str_name)

    # fidi.write('\n')

    print('Writing variables section of Dakota input file')

    fidi.write('variables, \n')

    for i in range(len(dmeth.variables)):
        str_name = dmeth.variables[i]
        # TODO: Remove this check after continuous_state.py has been updated!
        if str_name != 'continuous_state':
            var = eval('{}.{}()'.format(str_name, str_name))
            var.dakota_write(fidi, dvar)

    for i in range(len(dmeth.lcspec)):
        str_name = dmeth.lcspec[i]
        var = eval('{}.{}()'.format(str_name, str_name))
        # check that str_name is correct against matlab version which has no argument there
        var.dakota_write(fidi, eval('dvar.{}[i]'.format(str_name)), str_name)

    fidi.write('\n')


#  function to write the interface section of the file
def interface_write(fidi, params):

    print('Writing interface section of Dakota input file')

    fidi.write('interface, \n')

    if (not params.system) and (not params.fork) and (not params.direct):
        params.fork = True
    elif params.system + params.fork + params.direct > 1:
        raise RuntimeError('Too many interfaces selected')
    if params.system or params.fork:
        param_write(fidi, '\t', 'asynchronous', '', '\n', params)
        param_write(fidi, '\t  ', 'evaluation_concurrency', '=', '\n', params)
        param_write(fidi, '\t  ', 'analysis_concurrency', '=', '\n', params)
        param_write(fidi, '\t  ', 'evaluation_servers', '=', '\n', params)
        param_write(fidi, '\t  ', 'evaluation_self_scheduling', '', '\n', params)
        param_write(fidi, '\t  ', 'evaluation_static_scheduling', '', '\n', params)
        param_write(fidi, '\t  ', 'analysis_servers', '=', '\n', params)
        param_write(fidi, '\t  ', 'analysis_self_scheduling', '', '\n', params)
        param_write(fidi, '\t  ', 'analysis_static_scheduling', '', '\n', params)
        param_write(fidi, '\t', 'algebraic_mappings', '=', '\n', params)
        param_write(fidi, '\t', 'system', '', '\n', params)
        param_write(fidi, '\t', 'fork', '', '\n', params)
        param_write(fidi, '\t  ', 'analysis_driver', ' = \'', '\'\n', params)
        if len(params.input_filter) != 0:
            param_write(fidi, '\t  ', 'input_filter', '=', '\n', params)

        if len(params.output_filter) != 0:
            param_write(fidi, '\t  ', 'output_filter', '=', '\n', params)

        param_write(fidi, '\t  ', 'failure_capture', '   ', '\n', params)
        param_write(fidi, '\t  ', 'deactivate', '        ', '\n', params)
        param_write(fidi, '\t  ', 'parameters_file', ' = \'', '\'\n', params)
        param_write(fidi, '\t  ', 'results_file', ' = \'', '\'\n', params)
        param_write(fidi, '\t  ', 'verbatim', '', '\n', params)
        param_write(fidi, '\t  ', 'aprepro', '', '\n', params)
        param_write(fidi, '\t  ', 'file_tag', '', '\n', params)
        param_write(fidi, '\t  ', 'file_save', '', '\n', params)
    elif params.direct:
        #  Error: asynchronous capability not yet supported in direct interfaces.
        #  Update: it is now possible to run in parallel in direct interfaces.
        param_write(fidi, '\t', 'algebraic_mappings', '=', '\n', params)
        param_write(fidi, '\t', 'direct', '', '\n', params)
        param_write(fidi, '\t  ', 'analysis_driver', ' = \'', '\'\n', params)
        if float(IssmConfig('_DAKOTA_VERSION_')[0]) < 6:
            param_write(fidi, '\t  ', 'evaluation_static_scheduling', '', '\n', params)
        else:
            param_write(fidi, '\t  ', 'evaluation_scheduling', ' ', '\n', params)
            param_write(fidi, '\t  ', 'processors_per_evaluation', '=', '\n', params)
        if len(params.analysis_components) != 0:
            pathstr, name, ext = fileparts(params.analysis_components)
            if ext != '':
                ext = '.py'

            params.analysis_components = fullfile(pathstr, name + ext)
            param_write(fidi, '\t  ', 'analysis_components', ' = \'', '\'\n', params)

        if len(params.input_filter) != 0:
            param_write(fidi, '\t  ', 'input_filter', '=', '\n', params)

        if len(params.output_filter) != 0:
            param_write(fidi, '\t  ', 'output_filter', '=', '\n', params)

        param_write(fidi, '\t  ', 'failure_capture', '   ', '\n', params)
        param_write(fidi, '\t  ', 'deactivate', '        ', '\n', params)
        param_write(fidi, '\t  ', 'processors_per_analysis', '=', '\n', params)

    fidi.write('\n')


#  function to write the responses section of the file
def responses_write(fidi, dmeth, dresp, params):

    print('Writing responses section of Dakota input file')

    fidi.write('responses, \n')
    #fidi.write('calibration_terms = 1 \n')

    #  functions, gradients, and hessians vary by method

    rdesc = []
    for i in range(len(dmeth.responses)):
        resp = eval(dmeth.responses[i])
        rdesc = resp.dakota_write(fidi, dresp, rdesc)

    #  write accumulated response descriptors for all response classes

    if len(rdesc) != 0:
        fidi.write('\tresponse_descriptors =\n')
        vector_write(fidi, '\t  ', rdesc, 6, 76)

    ghspec_write(fidi, params, dmeth.ghspec)

    fidi.write('\n')


#  function to write gradient and hessian specifications
def ghspec_write(fidi, params, ghspec):

    #  gradients
    if 'grad' in ghspec:
        if (not params.numerical_gradients) and (not params.analytic_gradients):
            params.numerical_gradients = True
        elif (params.numerical_gradients + params.analytic_gradients > 1):
            raise RuntimeError('Too many gradients selected')

        if params.numerical_gradients:
            param_write(fidi, '\t', 'numerical_gradients', '', '\n', params)
            param_write(fidi, '\t  ', 'method_source', ' ', '\n', params)
            param_write(fidi, '\t  ', 'interval_type', ' ', '\n', params)
            param_write(fidi, '\t  ', 'fd_gradient_step_size', '=', '\n', params)
        elif params.analytic_gradients:
            param_write(fidi, '\t', 'analytic_gradients', '', '\n', params)
    #    elif params.mixed_gradients
    else:
        fidi.write('\tno_gradients\n')

    #  hessians (no implemented methods use them yet)
    if 'hess' in ghspec:
        raise RuntimeError('Hessians needed by method but not provided')
    else:
        fidi.write('\tno_hessians\n')
