import numpy as np
import os

class UcodeWriter(object):
    '''
    Class for writing Ucode input files

    Parameters
    ----------
    :OPTIONS BLOCK:
    verbose : int
        verbosity level (defalut is 3)

    :UCODE CONTROLS BLOCK:
    modelname : string
        Name of the model
    modellengthunits : string
        Length units of the model
    modelmassunits : string
        Mass units of the model
    modeltimeunits : string
        Time units of the model
    sensitivities : string
        Flag indicating whether sensitivities are calculated
    optimize : string
        Flag indicating whether optimization is performed
    linearity: string
        Flag indicating whether test model linearity mode is performed
    eigenvalues : string
        Flag indicating whether eigenvalues are printed
    startres : string
        Flag indicating whether starting residuals are printed
    intermedres : string
        Flag indicating whether intermediate residuals are printed
    finalres : string
        Flag indicating whether final residuals are printed
    dataexchange : string
        Flag indicating whether data exchange files are created

    :REG_GN_CONTROLS BLOCK:
    tolsosc : float
        Convergence criteria for objective function change
    maxiter : int
        Maximum number of parameter estimation iterations
    trustregion : string
        Flag indicating whether trustregion methods are used
    scaling : string
        Flag indicating whether parameter scaling is used

    :MODEL_COMMAND_LINES BLOCK:
    command : string
        string (2000 characters max) to execute the model(s)
    purpose : string
        Either forward or sensitivities
    commandid : string
        identifier for the command

    :PARAMETER_DATA_INPUT BLOCK:
    npar : int
        Number of parameters
    parname : string list of length npar
        The parameter name, such as 'HK_1'.
    startvalue : float list of length npar
        The starting value for the parameter.
    constrain: string list of length npar 
        Flag indicating whether constraints are applied to the parameter 
    lowerconstraint: float list of length npar 
        The lower bound constraint value for the parameter
    upperconstraint: float list of length npar 
        The upper bound constraint value for the parameter
    adjustable : string list of length npar
        Flag indicating whether parameter is adjustable
    perturbamt: float list of length npar 
        The fractional amount of a parameter value to perturb
    transform : string list of length npar
        Flag indicating whether parameter is log transformed
    tolpar : float list of length npar
        Convergence criteria for fractional parameter change per iteration
    maxchange : float list of length npar
        Maximum allowable fractional parameter change each iteration
    senmethod: int list of length npar
        Flag indicating how sensitivities are calculated

    :OBSERVATIONS_GROUP_INPUT BLOCK:
        NOT FULLY SUPPORTED.  HARDWIRED Still need to decide how to do this
    plotsymbol : integer list of length = to number of groups (likely same as nobsf)
        Integer value for plotsymbol. Used for post-processing
    useflag : string list of length = to number of groups (likely same as nobsf)
        Flag indicating if the group is used

    nobsf : int
        number of observation files
    nobs : int list of length nobsf
        number of observed values in each file

    :OBSERVATIONS_DATA_INPUT BLOCK:
    obsname : string list of length (nobsf,nobs)
        Name of the observation
    obsvalue : float list of length (nobsf,nobs)
        Value of the observation
    statistic : float list of length (nobsf,nobs)
        Value of the uncertainty statistic of the observation
    statflag : string list of length (nobsf,nobs)
        Statistic type of the observation (VAR, SD, CV, WT, or SQRWT)
    groupname : string list of length (nobsf,nobs)
        Name of the observation group

    :MODEL_INPUT_FILES BLOCK:
    ntplf : int
        number of template files
    modinfile : string list of length = number of template files
        path and file name of model input file
    templatefile : string list of length = number of template files
        path and file name of model template file

    :MODEL_OUTPUT_FILES BLOCK:
    modoutfile : string list of length = number of nobsf
        path and file name of model output file
    instructionfile : string list of length = number of nobsf
        path and file name of instruction file


    Attributes
    ----------

    Methods
    -------

    See Also
    --------

    Notes

    '''
    def __init__(self, model, verbose=3, modelname=[], modellengthunits=[],
                 modelmassunits=[], modeltimeunits=[], sensitivities='no', 
                 optimize='no', linearity='no',eigenvalues='yes', 
                 startres='yes', intermedres='no', finalres='yes', 
                 dataexchange='yes', tolsosc=0.0, maxiter=5,trustregion='no',
                 scaling='yes', command=[], purpose='forward', commandid=[], 
                 npar=0, parname=[], startvalue=[], constrain=[], 
                 lowerconstraint=[], upperconstraint=[], adjustable=[], 
                 perturbamt=[], transform=[], tolpar=[], maxchange=[], 
                 senmethod=[], nobsf=0, nobs=[], obsname=[], obsvalue=[], 
                 statistic=[], statflag=[], groupname=[], modinfile=[], 
                 templatefile=[], modoutfile=[], instructionfile=[]): 

        """ Package constructor
        """
        self.model = model
        self.verbose = verbose
        self.modelname = modelname
        self.modellengthunits = modellengthunits
        self.modelmassnits = modelmassunits
        self.modeltimeunits = modeltimeunits
        self.sensitivities = sensitivities
        self.optimize = optimize
        self.linearity = linearity
        self.eigenvalues = eigenvalues
        self.startres = startres
        self.intermedres = intermedres
        self.finalres = finalres
        self.dataexchange = dataexchange
        self.tolsosc = tolsosc
        self.maxiter = maxiter
        self.trustregion = trustregion
        self.scaling = scaling
        self.command = command
        self.purpose = purpose
        self.commandid = commandid
        self.npar = npar
        self.parname = parname
        self.startvalue = startvalue
        self.constrain = constrain
        self.lowerconstraint = lowerconstraint
        self.upperconstraint = upperconstraint
        self.adjustable = adjustable
        self.perturbamt = perturbamt
        self.transform = transform
        self.tolpar = tolpar
        self.maxchange = maxchange
        self.senmethod = senmethod
        self.nobsf = nobsf
        self.nobs = nobs
        self.obsname = obsname
        self.obsvalue = obsvalue
        self.statistic = statistic
        self.statflag = statflag
        self.groupname = groupname
        self.modinfile = modinfile
        self.templatefile = templatefile
        self.modoutfile = modoutfile
        self.instructionfile = instructionfile


        # -create empty arrays of the correct size
        mxnobs = max(self.nobs) 
#swm:debug        print 'nobsf and nobs', self.nobsf, max(self.nobs)
        self.startvalue = np.zeros((self.npar), dtype='float32')
#        self.obsvalue = np.zeros((self.nobsf, max(self.nobs)), dtype='float32')
        #swm: might not need to assign obsvalue - obsname (5 lines below)
        self.obsvalue = np.zeros((self.nobsf, mxnobs), dtype='float32')
        self.statistic = np.zeros((self.nobsf, mxnobs), dtype='float32')
        self.statflag = np.zeros((self.nobsf, mxnobs), dtype=object)
        self.groupname = np.zeros((self.nobsf, mxnobs), dtype=object)
        self.obsname = np.zeros((self.nobsf, mxnobs), dtype=object)
        self.nobs = np.zeros((self.nobsf), dtype='int32')

        # -assign values to arrays
        # obsname thru groupname are variable length lists, so padding with None
        length = len(sorted(obsname,key=len, reverse=True)[0])
        self.obsname=np.array([xi+[None]*(length-len(xi)) for xi in obsname])
        self.obsvalue=np.array([xi+[None]*(length-len(xi)) for xi in obsvalue])
        self.statistic=np.array([xi+[None]*(length-len(xi)) for xi in statistic])
        self.statflag=np.array([xi+[None]*(length-len(xi)) for xi in statflag])
        self.groupname=np.array([xi+[None]*(length-len(xi)) for xi in groupname])

        self.parname[:] = parname
        self.startvalue[:] = startvalue
        self.constrain[:] = constrain
        self.lowerconstraint[:] = lowerconstraint
        self.upperconstraint[:] = upperconstraint
        self.adjustable[:] = adjustable
        self.perturbamt[:] = perturbamt
        self.transform[:] = transform
        self.tolpar[:] = tolpar
        self.maxchange[:] = maxchange
        self.senmethod[:] = senmethod
        
        self.nobs[:] = nobs
        self.modinfile[:] = modinfile
        self.templatefile[:] = templatefile
        self.modoutfile[:] = modoutfile
        self.instructionfile[:] = instructionfile

        # putting in some more checks here


        # add checks for input compliance (obsnam length, etc.)

    def write_file(self):
        """
        Write the ucode input file

        Returns
        -------
        None

        """
        # -open file for writing
#        i = open(self.modelname + '_ucode.in', 'w') 
#        i = open(self.model.model_ws + '\\' + self.modelname + '_ucode.in', 'w')
        i = open(os.path.join(self.model.model_ws, self.modelname + '_ucode.in'), 'w')
        separ = '#'+'-'*40
       ## UCODE block ##
        i.write("BEGIN Options  KEYWORDS\n")
        i.write("  Verbose=%s\n" %self.verbose)
        i.write("END Options\n\n")

#        ## UCODE-CONTROL INFORMATION block ##
        i.write("%s\n# UCODE-CONTROL INFORMATION \n%s\n\n"%(separ,separ))
        i.write("BEGIN UCODE_CONTROL_DATA KEYWORDS\n")
        i.write("ModelName=%s\n"%self.modelname)
        i.write("sensitivities=%s\n"%self.sensitivities)
        i.write("optimize=%s\n"%self.optimize)
        i.write("linearity=%s\n"%self.linearity)
        i.write("eigenvalues=%s\n"%self.eigenvalues)
        i.write("startres=%s\n"%self.startres)
        i.write("intermedres=%s\n"%self.intermedres)
        i.write("finalres=%s\n"%self.finalres)
        i.write("dataexchange=%s\n"%self.dataexchange)
#        i.write("OmitInsensitive=%s\n"%omit)
        i.write("END UCODE_CONTROL_DATA\n\n")
        i.write("BEGIN REG_GN_CONTROLS KEYWORDS\n")
#        i.write("tolpar=%s\n"%self.tolpar)
        i.write("tolsosc=%s\n"%self.tolsosc)
        i.write("maxiter=%s\n"%self.maxiter)
#        i.write("maxchange=%s\n"%self.maxchange)
#        i.write("MaxchangeRealm=%s\n"%maxchange_realm)
        i.write("trustregion=%s\n"%self.trustregion)
        i.write("scaling=%s\n"%self.scaling)
        i.write("END REG_GN_CONTROLS\n\n")

        ## COMMAND FOR APPLICATION MODEL(S) block ##
        i.write("%s\n# COMMAND FOR APPLICATION MODEL(S) \n%s\n\n"%(separ,separ))
        i.write("BEGIN MODEL_COMMAND_LINES TABLE\n")
        i.write("nrow=1 ncol=3  columnlabels\n")
        i.write("COMMAND    PURPOSE    COMMANDID\n")
        i.write("%s %s  %s\n"%(self.command,self.purpose,self.commandid))
        i.write("END MODEL_COMMAND_LINES\n\n")

        i.write("%s\n# PARAMETER INFORMATION \n%s\n\n"%(separ,separ))
        i.write("BEGIN PARAMETER_DATA TABLE\n")
        i.write("nrow=%s ncol=%s columnlabels\n"%(self.npar,11))    #swm: note ncols hardwired!!!
        i.write("ParamName  StartValue  Constrain  LowerConstraint  UpperConstraint  Adjustable  PerturbAmt  Transform  TolPar  MaxChange  SenMethod\n")    #swm: note ncols hardwired!!!

#swm: Note: formatting for adding spaces seems a bit crude.  Probably a better way to do this
        for np in range(0,self.npar):
            i.write(
             '{}{:10.4g}{}{:10.4g}{}{:10.4g}{}{:10.4g}{}{:10.4g}{}{:10.4g}{}{:10d}\n'
              .format(self.parname[np] + ' ', 
               self.startvalue[np], 
               ' '*10 + self.constrain[np] + ' '*10, 
               self.lowerconstraint[np], 
               ' ', self.upperconstraint[np], 
               ' '*15 + self.adjustable[np] + ' ', 
               self.perturbamt[np], 
               ' '*5 + self.transform[np] + ' ',
               self.tolpar[np],
               ' ', self.maxchange[np],
               ' ', self.senmethod[np]))
        i.write("END PARAMETER_DATA\n\n")

        i.write("%s\n# OBSERVATION INFORMATION \n%s\n\n"%(separ,separ))

        i.write("BEGIN OBSERVATION_GROUPS TABLE\n")
        i.write("NROW=%s  NCOL=%s  COLUMNLABELS\n"%((self.nobsf),3)) #swm: hardwired cols
        i.write("GroupName  PlotSymbol  UseFlag\n")    #swm: note ncols hardwired!!!
        for j in range(0,self.nobsf):
                i.write('{}{}{:10d}{}\n'.format(self.groupname[j,0] , ' ',
                                                   j+1, ' yes'))
        i.write("END OBSERVATION_GROUPS\n\n")

        i.write("BEGIN OBSERVATION_DATA TABLE\n")
        i.write("NROW=%s  NCOL=%s  COLUMNLABELS\n"%((sum(self.nobs)),5)) #swm: hardwired cols
        i.write("ObsName  ObsValue  Statistic StatFlag  GroupName\n")    #swm: note ncols hardwired!!!

        for j in range(0,self.nobsf):
            for k in range(0,self.nobs[j]):
#                i.write('{}{:10.4g}{:10.4g}{}{}\n'.format(self.obsname[j,k] + ' ', #swm: not sure why this format statement doesn't work
                i.write('{}{}{}{}{}{}{}{}{}\n'.format(self.obsname[j,k],
                                            ' ', self.obsvalue[j,k],
                                            ' ', self.statistic[j,k],
                                            ' ', self.statflag[j,k],
                                            ' ', self.groupname[j,k]))
        i.write("END OBSERVATION_DATA\n\n")

        i.write("%s\n# APPLICATION MODEL INFORMATION \n%s\n\n"%(separ,separ))
        i.write("BEGIN MODEL_INPUT_FILES        KEYWORDS\n")
        for nf in range(len(self.templatefile)):     # swm: is this right?  Or do I need to specify?
            i.write("  modinfile=%s\n"%self.modinfile[nf])
            i.write("  templatefile=%s\n"%self.templatefile[nf])
        i.write("END MODEL_INPUT_FILES\n\n")

        i.write("BEGIN MODEL_OUTPUT_FILES  KEYWORDS\n")
        for nf in range(0, self.nobsf):
            i.write("  modoutfile=%s\n"%self.modoutfile[nf])
            i.write("  instructionfile=%s\n"%self.instructionfile[nf])
            i.write("  category=obs\n")
        i.write("END MODEL_OUTPUT_FILES")

        i.close()


        return
