"""
Classes for writing Pipeline input files.
Copyright Schrodinger, LLC. All rights reserved.
"""
# Contributors: Matvey Adzhigirey
import os
from schrodinger.application import inputconfig
[docs]class Writer:
    """
    Class for writing pipeline input files.
    """
[docs]    def __init__(self, filename=None, comment=None):
        self.config = inputconfig.InputConfig()
        self.filename = filename
        if comment:
            self.config.initial_comment = comment.splitlines()
        self.userouts = set()
        self.structout = None
        self._used_names = [] 
[docs]    def write(self, filename=None):
        """
        Serialize the job parameters to a file.
        """
        if filename is None:
            filename = self.filename
        if not filename.endswith('.inp'):
            raise ValueError(
                "VSW.input.write(): File name must have *.inp extension")
        # Write the USEROUTS section (if any):
        if self.userouts or self.structout:
            section = "USEROUTS"
            self.config[section] = {}
            if self.userouts:
                self.config[section]["USEROUTS"] = list(self.userouts)
            if self.structout:
                self.config[section]["STRUCTOUT"] = self.structout
        self.config.writeInputFile(filename, yesno=True)
        return filename 
[docs]    def addVar(self, vname, vclass, data, comment=None):
        """
        Add a variable of name `vname` and class `vclass`. For `vclass`
        of "PhaseDB" and "Grid", `data` should be a single path; for
        `vclass` of "Structures", `data` should be a list of files.
        The argument `comment` is ignored.
        :type vname: str
        :type vclass: str
        """
        if vname in self._used_names:
            raise ValueError("Name '%s' already in use" % vname)
        self._used_names.append(vname)
        section = "SET:%s" % vname
        self.config[section] = {}
        self.config[section]["VARCLASS"] = vclass
        if vclass == 'PhaseDB':
            self.config[section]["PATH"] = os.path.abspath(data)
        elif vclass == 'Grid':
            self.config[section]["FILE"] = os.path.abspath(data)
        else:
            self.config[section]["FILES"] = [
                os.path.abspath(filename) for filename in data
            ] 
[docs]    def userOutput(self, varname):
        """
        Append output parameter descriptors to `self.userouts`.
        """
        # Make the output available to the user at the end.
        self.userouts.add(varname) 
[docs]    def setStructureOutput(self, varname):
        """
        Set structure output to `varname` (overwriting the previous one, if
        any). The structures from this variable will be imported into
        Maestro when incorporating.
        """
        self.structout = varname 
[docs]    def addStage(self,
                 sname,
                 sclass,
                 inputs,
                 outputs,
                 keywords=None,
                 comment=None):
        """
        Add a stage.
        The `comment` argument is ignored.
        """
        if sname in self._used_names:
            raise ValueError("Name '%s' already in use" % sname)
        self._used_names.append(sname)
        # keywords - list of tuples to retain the specified order
        section = "STAGE:%s" % sname
        self.config[section] = {}
        if not keywords:
            keywords = []
        # Convert dictionary to a list of tuples:
        if isinstance(keywords, type({})):
            kw = []
            for key, value in keywords.items():
                kw.append((key, value))
        else:
            kw = keywords
        self.config[section]["STAGECLASS"] = sclass
        if inputs:
            self.config[section]["INPUTS"] = inputs
        if outputs:
            self.config[section]["OUTPUTS"] = outputs
        # Save the sections, as they need to be added at the very bottom:
        sections = []
        for k in kw:
            key = k[0]
            value = k[1]
            if isinstance(value, type(True)):
                # Is a boolean type
                if value:
                    value = 'YES'
                else:
                    value = 'NO'
            if isinstance(value, type({})):
                # Save the section to have it printed at the bottom of the
                # file:
                sections.append((key, value))
            else:
                self.config[section][key] = value
        for key, section_dict in sections:
            self.config[section][key] = section_dict  
[docs]class LinearWriter(Writer):
    """
    Class for writing linear stage maps (ones with no branches).
    """
[docs]    def __init__(self, filename, comment=None):
        Writer.__init__(self, filename, comment)
        self.last_output = None 
[docs]    def addVar(self, vname, vclass, files, comment=None):
        """
        """
        Writer.addVar(self, vname, vclass, files, comment)
        self.last_output = vname 
[docs]    def addStage(self, sname, sclass, output, keywords=None, comment=None):
        """
        """
        Writer.addStage(self, sname, sclass, [self.last_output], [output],
                        keywords, comment)
        self.last_output = output  
# EOF