from CCP4PluginScript import CPluginScript

class gesamt(CPluginScript):

    TASKTITLE='Gesamt - structural alignment'
    TASKNAME = 'gesamt'
    TASKMODULE= 'model_data_utility'
    TASKCOMMAND = 'gesamt'
    TASKVERSION= 0.0
    PERFORMANCECLASS = 'CSuperposePerformance'

    def makeCommandAndScript(self):

      inp = self.container.inputData
      par = self.container.controlParameters
      out = self.container.outputData

      import CCP4Utils
      import os

      self.appendCommandLine( [ str( inp.XYZIN_QUERY.fullPath ) ] )
      if inp.XYZIN_QUERY.isSelectionSet():
        self.appendCommandLine( [ '-s' , inp.XYZIN_QUERY.selection.__str__() ] )
      if inp.XYZIN_TARGET.fullPath.isSet():
        self.appendCommandLine( [ str( inp.XYZIN_TARGET.fullPath ) ] )
      if inp.XYZIN_TARGET.isSelectionSet():
        self.appendCommandLine( [ '-s' ,  inp.XYZIN_TARGET.selection.__str__()  ] )
      if par.MODE.isSet():
        self.appendCommandLine( [ '-'+str( par.MODE ) ] )

      if out.XYZOUT.fullPath.isSet():
        if par.OUTPUT_COORDS == "minusO":
          self.appendCommandLine( [ '-o', str( out.XYZOUT.fullPath ) ] )
        elif par.OUTPUT_COORDS == "minusOF":
          a,b = os.path.splitext ( inp.XYZIN_TARGET.fullPath )
          self.tmpOut = a + "_2" + b
          self.appendCommandLine( [ '-o-f' ] )

    def processOutputFiles(self):
        logName = self.makeFileName('LOG')
        eulerValues = None
        translationValues = None
        rmsValue = None
        qValue = None
        nResValue = None
        transformationMatrix = []
        
        from lxml import etree
        xmlRoot = etree.Element('Gesamt')
        perResidueNode = None
        iRes = 0

        import os
        out = self.container.outputData
        if hasattr(self,"tmpOut") and os.path.isfile(self.tmpOut):
          os.rename ( self.tmpOut,str(out.XYZOUT.fullPath) )

        with open (logName,'r') as logFile:
            lines = logFile.readlines()
            inPerResidue = False
            inQueryMatrix = False
            inMeatOfMatrix = False
            for line in lines:
                tokens = line.strip().split()
                if line.strip().startswith('Euler angles (alpha,beta,gamma)'):
                    eulerValues = [float(tokens[-3]), float(tokens[-2]), float(tokens[-1])]
                if line.strip().startswith('Orthogonal translation (Angstrom)'):
                    translationValues = [float(tokens[-3]), float(tokens[-2]), float(tokens[-1])]
                if line.strip().startswith('RMSD'):
                    rmsValue = float(tokens[-1])
                if line.strip().startswith('Q-score'):
                    qValue = float(tokens[-1])
                if line.strip().startswith('Aligned residues'):
                    nResValue = int(tokens[-1])
                if line.strip().startswith('.-------------.------------.-------------'):
                    inPerResidue = True
                    perResidueNode = etree.Element('PerResidue')
                if 'Transformation matrix for Target:' in line: inQueryMatrix = True
                if inQueryMatrix:
                    if inMeatOfMatrix:
                        if len(line.strip()) == 0:
                            inQueryMatrix = False
                            inMeatOfMatrix = False
                        else:
                            for value in line.strip().split():
                                transformationMatrix.append(float(value))
                    if 'Rx' in line:
                        inMeatOfMatrix = True
                if inPerResidue:
                    tokens = line.strip().split('|')
                    if len(tokens) > 1 and not ('Query' in tokens[1] or '-----' in tokens[1]):
                        row = etree.SubElement(perResidueNode,'Row')
                        calphaDistanceNode = etree.SubElement(row,'Distance')
                        calphaDistanceNode.text = tokens[2].strip()[3:-3]
                        equivalenceNode = etree.SubElement(row,'Equivalence')
                        equivalenceNode.text = line.strip()
                        iResNode = etree.SubElement(row,'iRes')
                        iResNode.text = str(iRes)
                        iRes += 1
                if line.strip().startswith("`-------------'------------'-------------'"):
                    inPerResidue = False

    
        if eulerValues is not None and translationValues is not None and rmsValue is not None and qValue is not None and nResValue is not None:
            from CCP4MathsData import CTransformation
            self.container.outputData.TRANSFORMATION = CTransformation()
            self.container.outputData.TRANSFORMATION.alpha = eulerValues[0]
            self.container.outputData.TRANSFORMATION.beta = eulerValues[1]
            self.container.outputData.TRANSFORMATION.gamma = eulerValues[2]
            self.container.outputData.TRANSFORMATION.x = translationValues[0]
            self.container.outputData.TRANSFORMATION.y = translationValues[1]
            self.container.outputData.TRANSFORMATION.z = translationValues[2]
            self.container.outputData.PERFORMANCE.RMSxyz = rmsValue
            #self.container.outputData.PERFORMANCE.QScore = qValue
            self.container.outputData.PERFORMANCE.nResidues = nResValue
            transformationNode = etree.SubElement(xmlRoot,'Transformation',
                                                  alpha=str(self.container.outputData.TRANSFORMATION.alpha),
                                                  beta=str(self.container.outputData.TRANSFORMATION.beta),
                                                  gamma=str(self.container.outputData.TRANSFORMATION.gamma),
                                                  x = str(self.container.outputData.TRANSFORMATION.x ),
                                                  y = str(self.container.outputData.TRANSFORMATION.y ),
                                                  z = str(self.container.outputData.TRANSFORMATION.z),
                                                  rms = str(rmsValue),
                                                  q = str( qValue ),
                                                  nRes  = str(nResValue ) )
            matrixNode = etree.SubElement(transformationNode,'Matrix')
            
            #Storing the transform in 4x4 matrix format, we need to add the bottom row of a 4x4transformation matrix
            for element in transformationMatrix + [0.,0.,0.,1.]:
                elementNode = etree.SubElement(matrixNode,'Element')
                elementNode.text = str(element)

            if perResidueNode is not None: transformationNode.append(perResidueNode)
        
        with open(self.makeFileName('PROGRAMXML'),'w') as xmlFile:
            xmlString = etree.tostring(xmlRoot, pretty_print=True)
            xmlFile.write(xmlString)

        return CPluginScript.SUCCEEDED