"""
    import_merged.py: CCP4 GUI Project
     Copyright (C) 2015 STFC

     This library is free software: you can redistribute it and/or
     modify it under the terms of the GNU Lesser General Public License
     version 3, modified in accordance with the provisions of the 
     license to address the requirements of UK law.
 
     You should have received a copy of the modified GNU Lesser General 
     Public License along with this library.  If not, copies may be 
     downloaded from http://www.ccp4.ac.uk/ccp4license.php
 
     This program is distributed in the hope that it will be useful,
     but WITHOUT ANY WARRANTY; without even the implied warranty of
     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
     GNU Lesser General Public License for more details.
"""

import os,shutil
from CCP4PluginScript import CPluginScript
import CCP4Utils
from CCP4ErrorHandling import *
from lxml import etree

class import_merged(CPluginScript):

    TASKNAME = 'import_merged'
    MAINTAINER = 'liz.potterton@york.ac.uk'
    WHATNEXT = []
    ERROR_CODES = { 301 : { 'description' : 'No output file found after conversion program' },
                    302 : { 'description' : 'Output from conversion does not contain recognised reflection or FreeR set data' }
                    }
    # Note - preserving the HKLOUT by changing severity from the systm default of 1 to 5 and
    # beware issues ith caseinsensitivity
    PURGESEARCHLIST = [ [ 'HKLIN*.mtz' , 1 ],
                        ['aimless_pipe%*/HKLOUT*.mtz', 1],
                         [ 'hklout.mtz' , 5 ],    
                         [ 'HKLOUT.mtz' , 5 ]
                       ]
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def process(self):
      self.container.inputData.HKLIN.loadFile()
      self.fformat = self.container.inputData.HKLIN.getFormat()
      # Type(self.fformat) can be either [mtz] <class 'CCP4Data.CString'> or [mmcif] str  WHY? 
      print "self.fformat", type(self.fformat)

      merged = self.container.inputData.HKLIN.getMerged()
      self.isintensity = 0  # unknown I or F
      
      print 'Imported file of format:',self.fformat,merged
      print 'HKLIN_OBS_COLUMNS',self.container.inputData.HKLIN_OBS_COLUMNS,type(self.container.inputData.HKLIN_OBS_COLUMNS),self.container.inputData.HKLIN_OBS_COLUMNS.isSet()
      obsout = None

      self.x2mtz = None
      if self.fformat == 'mtz':
        self.x2mtz = self.makePluginObject('x2mtz')
      elif self.fformat == 'mmcif':
        self.x2mtz = self.makePluginObject('cif2mtz')
      elif self.fformat in [ 'shelx' ]:
        self.x2mtz = self.makePluginObject('convert2mtz')
      elif self.fformat in [ 'sca' ]:
        self.x2mtz = self.makePluginObject('scalepack2mtz')

      self.x2mtz.container.inputData.copyData(otherContainer=self.container.inputData)
      self.x2mtz.container.outputData.copyData(otherContainer=self.container.outputData,dataList=['HKLOUT','OBSOUT'])


      self.freeRcompleteTried = True
      if self.fformat == 'mtz':
        # Just call the processOutputFiles() to convert to mini mtzs
        self.x2mtz.container.outputData.HKLOUT.set(self.container.inputData.HKLIN)
        # Pick up dataset name (and crystal name if available)
        fcontent = self.container.inputData.HKLIN.getFileContent()
        print 'input file content', type(fcontent), fcontent
        print 'columns', fcontent.listOfColumns
        # check if selection columns are intensity or amplitudes
        # return +1 if intensity, -1 if amplitude, 0 if unknown
        self.isintensity = self.isIntensity(self.container.inputData.HKLIN_OBS_COLUMNS,
                                            fcontent.listOfColumns)
        self.importXML = etree.Element('IMPORT_LOG')  # information about the import step
        self.makeReportXML(self.importXML)  # add initial stuff for XML into self.importXML
        self.outputLogXML(self.importXML)  # send self.importXML to program.xml

        if len(fcontent.datasets)>=2:
          #print fcontent.datasets[1]
          # self.crystalName = ###  set this, but how?
          self.container.inputData.DATASETNAME = fcontent.datasets[1]

        self.x2mtz.checkOutputData()
        ret = self.x2mtz.processOutputFiles()  # runs cmtzsplit
        self.x2mtz.reportStatus(ret)
        self.container.outputData.OBSOUT.set(self.x2mtz.container.outputData.OBSOUT)
        print 'OBSOUT',self.container.outputData.OBSOUT.get()
        self.process1({'finishStatus': ret })
      else:
          self.importXML = etree.Element('IMPORT_LOG')  # information about the import step
          self.makeReportXML(self.importXML)  # add initial stuff for XML into self.importXML
          self.outputLogXML(self.importXML)  # send self.importXML to program.xml
          self.connectSignal(self.x2mtz,'finished',self.process1)
          ret = self.x2mtz.process()
        
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def process1(self,status, completeFreeR=True):
        'if completeFreeR False, always generate new FreeR (for 2nd attempt)'
        print 'process1',status
        if status is not None and status.get('finishStatus') == CPluginScript.FAILED:
            self.reportStatus(status)
            return
      
        # Create or complete a freer set
        self.freerflag = self.makePluginObject('freerflag')
        self.freerflag.container.inputData.F_SIGF = self.x2mtz.container.outputData.OBSOUT
        print 'import_merged.process1',self.x2mtz.container.outputData.FREEOUT,self.x2mtz.container.outputData.FREEOUT.exists()
        newfreer = 'True'
        if completeFreeR and self.x2mtz.container.outputData.FREEOUT.exists():
            # A freeR set has been imported, so extend/complete it
            self.freerflag.container.inputData.FREERFLAG = self.x2mtz.container.outputData.FREEOUT
            self.freerflag.container.controlParameters.GEN_MODE = 'COMPLETE'
            self.freerflag.container.controlParameters.COMPLETE = True
            newfreer = 'False'
          
        self.addElement(self.importXML, 'newFreeR', newfreer) 
        self.outputLogXML(self.importXML)  # send self.importXML to program.xml
        self.freerflag.container.outputData.FREEROUT.setFullPath(str(self.container.outputData.FREEOUT))
        self.connectSignal(self.freerflag,'finished',self.process2)
        status = self.freerflag.process()
        
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def process2(self,status):
      print 'in process 2, process1',status
      freerOK = True
      if status is not None and status.get('finishStatus') == CPluginScript.FAILED:
          # FreeR run has failed, create error message and continue
          print "process2 Failed"
          self.addElement(self.importXML, 'FreeRfailed', 'True')
          self.outputLogXML(self.importXML)  # send self.importXML to program.xml
          # try again
          if self.freeRcompleteTried:
              print "trying again"
              self.freeRcompleteTried = True
              self.process1(None, False)  # try to make a new FreeR set
          else:
              # failed a 2nd time, make error and continue
              print "failed again"
              self.addElement(self.importXML, 'FreeRfailed', 'Again')
              self.outputLogXML(self.importXML)  # send self.importXML to program.xml
              freerOK = False

          # self.reportStatus(status)

          # return

      # If FreeR is OK, then create annotation
      if freerOK:
          print "FREEROUT content",self.freerflag.container.outputData.FREEROUT.fileContent
          if self.freerflag.container.outputData.FREEROUT.fileContent.spaceGroup.isSet():
              sgname = self.freerflag.container.outputData.FREEROUT.fileContent.spaceGroup.__str__()
          else:
              sgname = 'Unk'

          highresFRformatted = "%7.2f" % float(self.freerflag.container.outputData.FREEROUT.fileContent.resolutionRange.high)
          title ='FreeR - Spg:'+str(sgname).strip()+';Resln:'+highresFRformatted.strip() + "A;"
          try:
              title = title + "Cell:"+self.freerflag.container.outputData.FREEROUT.fileContent.cell.guiLabel()
          except Exception as e:
              print 'Error writing cell parameters',e

          print "FreeR title:",title
          self.container.outputData.FREEOUT.annotation = title

      # Run aimless for a report on data quality
      self.aimless = self.makePluginObject('aimless_pipe',pluginTitle='DR run for data analysis')
      unmergedList = self.aimless.container.inputData.UNMERGEDFILES
      #print '\nunmergedList 0',    unmergedList
      if len(unmergedList)==0: unmergedList.addItem()
      if self.fformat in ['mmcif']:
        unmergedList[0].file.set(self.container.outputData.HKLOUT.__str__())
      elif  self.fformat in ['mtz']:
        unmergedList[0].file.set(self.container.outputData.OBSOUT.__str__())
      else:
        unmergedList[0].file.set(self.container.inputData.HKLIN.__str__())

      print '\nprocess1 self.container.inputData', self.container.inputData, '\n'
      print 'unmergedList 1',    unmergedList
      unmergedList[0].crystalName.set(self.container.inputData.CRYSTALNAME)
      unmergedList[0].dataset.set(self.container.inputData.DATASETNAME)
      #print 'unmergedList 2',    unmergedList
      unmergedList[0].cell.set(self.container.inputData.SPACEGROUPCELL.cell)
      #print 'unmergedList 3',    unmergedList
      print 'self.container.inputData',self.container.inputData
      print 'self.container.inputData.SPACEGROUPCELL',\
            self.container.inputData.SPACEGROUPCELL
      unmergedList[0].wavelength.set(self.container.inputData.WAVELENGTH)
      print 'unmergedList 4',    unmergedList, '\n'

      # parameters for Pointless
      self.aimless.container.controlParameters.MODE = 'CHOOSE'
      self.aimless.container.controlParameters.CHOOSE_MODE = 'SPACEGROUP'
      self.aimless.container.controlParameters.CHOOSE_SPACEGROUP = \
            self.container.inputData.SPACEGROUPCELL.spaceGroup
      # parameters for Aimless
      self.aimless.container.controlParameters.SCALING_PROTOCOL = 'CONSTANT'
      self.aimless.container.controlParameters.ONLYMERGE = True
      self.aimless.container.controlParameters.ANALYSIS_MODE = True
      self.aimless.container.controlParameters.SDCORRECTION_OVERRIDE = True
      self.aimless.container.controlParameters.SDCORRECTION_REFINE = False
      tempXML = self.importXML
      self.addElement(tempXML, "DRPIPE_RUNNING", "True") 
      self.outputLogXML(tempXML)  # SEND self.importXML to program.xml
      #  Start data reduction
      self.connectSignal(self.aimless,'finished',self.nearlyDone)
      self.aimless.process()

      return CPluginScript.SUCCEEDED

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def nearlyDone(self,status):
      print 'import_merged.nearlyDone'
      self.container.outputData.OBSOUT.setContentFlag(reset=True)
      import shutil
      try:
          # XML data: We have
          #   a) self.importXML etree element report on the import step
          #   b) aimless pipe report
          # so put these together into an IMPORT_MERGED block
          aimless_pipe_XMLpath = self.aimless.makeFileName('PROGRAMXML')
          aimless_pipe_Etree = etree.parse(aimless_pipe_XMLpath)
          aimless_pipe = aimless_pipe_Etree.getroot()
          print 'aimless_pipe', type(aimless_pipe), aimless_pipe
          self.outputLogXML(self.importXML, aimless_pipe)  # and output it
      except:
        pass
      self.reportStatus(status)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def isIntensity(self, selectedcolumns, listOfColumns):
        """
        check if selection columns are intensity or amplitudes
        return +1 if intensity, -1 if amplitude, 0 if unknown
        selectedcolumns  string of columns labels
        listOfColumns  list eg
        [{'groupIndex': '1', 'columnLabel': 'F_xe1a', 'columnType': 'F', 'dataset': 'xe1a'},
        """
        print 'isIntensity',selectedcolumns,type(selectedcolumns)
        if not selectedcolumns.isSet():
          selcol1 = str(listOfColumns[0].get()['columnLabel'])
        else:
          selected = selectedcolumns.split(',')
          selcol1 = selected[0] # first label
        print 'selected', selectedcolumns, selcol1
        columntype = None
        for col in listOfColumns:
            column = col.get()
            if column['columnLabel'] == selcol1:
                print 'found', column['columnLabel'], column['columnType']
                columntype = column['columnType']
        isintensity = 0
        if columntype is None:
            print "Unrecognised column " + selcol1
        elif (columntype == 'J') or (columntype == 'K'):
            isintensity = +1
        elif (columntype == 'F') or (columntype == 'G'):
            isintensity = -1

        return isintensity


    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def outputLogXML(self, x1XML, x2XML=None):
      'output x1XML and optionally x2XML to program.xml'
      rootXML = etree.Element('IMPORT_MERGED') # Global XML for everything
      rootXML.append(x1XML)
      if x2XML is not None:
          rootXML.append(x2XML)
      print "writing rootXML"
      with open (self.makeFileName('PROGRAMXML'),"w") as outputXML:
          outputXML.write(etree.tostring(rootXML,pretty_print=True))


    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def makeReportXML(self, containerXML):
        'Make initial report XML'
        print "mrx 1",  type(self.fformat),  self.fformat
        self.addElement(containerXML, 'fileformat', str(self.fformat))
        print 'type merged', type(self.container.inputData.HKLIN.getMerged())
        if self.container.inputData.HKLIN.getMerged():
            self.addElement(containerXML, 'merged', 'True')
        else:
            self.addElement(containerXML, 'merged', 'False')
        if self.fformat == 'mtz':
            self.addElement(containerXML, 'columnlabels',
                            self.container.inputData.HKLIN_OBS_COLUMNS.get())
        # = +1 if intensity, -1 if amplitude, 0 if unknown
        if self.isintensity == 0:
            IorFtype = 'Unknown'
        elif self.isintensity > 0:
            IorFtype = 'Intensity'
        elif self.isintensity < 0:
            IorFtype = 'Amplitude'
        self.addElement(containerXML, 'IorFtype', IorFtype)
        
            
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 
    def addElement(self, containerXML, elementname, elementtext):
        #print 'addElement', elementname, type(elementtext), elementtext 
        e2 = etree.Element(elementname)
        e2.text = elementtext
        containerXML.append(e2)

#----------------------------------------------------------------------------------------------------

"""
# Function to return list of names of exportable MTZ(s)
def exportJobFile(jobId=None,mode=None):
    import os
    import CCP4Modules
    import CCP4XtalData

    jobDir = CCP4Modules.PROJECTSMANAGER().jobDirectory(jobId=jobId,create=False)
    exportFile = os.path.join(jobDir,'exportMtz.mtz')
    if os.path.exists(exportFile): return exportFile

    childJobs = CCP4Modules.PROJECTSMANAGER().db().getChildJobs(jobId=jobId,details=True)
    print 'import_merged.exportMtz',childJobs
    truncateOut = None
    freerflagOut = None
    for jobNo,subJobId,taskName  in childJobs:
      if taskName == 'aimless_pipe':
         aimlessChildJobs = CCP4Modules.PROJECTSMANAGER().db().getChildJobs(jobId=subJobId,details=True)
         print 'import_merged.exportMtz aimlessChildJobs',aimlessChildJobs
         for jobNo0,subJobId0,taskName0  in aimlessChildJobs:
           if taskName0 == 'ctruncate':
             truncateOut = os.path.join( CCP4Modules.PROJECTSMANAGER().jobDirectory(jobId=subJobId0,create=False),'HKLOUT.mtz')
             if not os.path.exists(truncateOut): truncateOut = None
     
    freerflagOut = os.path.join( CCP4Modules.PROJECTSMANAGER().jobDirectory(jobId=jobId,create=False),'FREEOUT.mtz')
    if not os.path.exists(freerflagOut): freerflagOut = None
    if truncateOut is None: return None
    if freerflagOut is None: return truncateOut

    print 'aimless_pipe.exportJobFile  runCad:',exportFile,[ freerflagOut ]
    

    m = CCP4XtalData.CMtzDataFile(truncateOut)
    #print m.runCad.__doc__   #Print out docs for the function
    outfile,err = m.runCad(exportFile,[ freerflagOut ] )
    print 'aimless_pipe.exportJobFile',outfile,err.report()
    return   outfile                                                   
 
def exportJobFileMenu(jobId=None):
    # Return a list of items to appear on the 'Export' menu - each has three subitems:
    # [ unique identifier - will be mode argument to exportJobFile() , menu item , mime type (see CCP4CustomMimeTypes module) ]
    return [ [ 'complete_mtz' ,'MTZ file' , 'application/CCP4-mtz' ] ]
                                                
"""