"""
     arcimboldo.py: CCP4 GUI 2 Project
     Copyright (C) 2014 The University of York

     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.
"""

## @package arcimboldo
# This script runs all three versions of arcimboldo

#!/usr/bin/env ccp4-python
import pyrvapi
import xml.etree.ElementTree as ET

import os, subprocess, sys, time, json, re
from core.CCP4PluginScript import CPluginScript
from core import CCP4XtalData
from core import CCP4Utils
from core import CCP4Modules

ccp4_home = os.environ.get ( "CCP4", "not_set" )

class arcimboldo_rvapi(object):

   def __init__(self):
      return

   def generate_report(self, wrk_path, fname, mode=None):
      file_path = os.path.join(wrk_path,fname)
      file_xml = file_path+'.xml'
      file_html = file_path+'.html'

      if not os.path.isdir(wrk_path) or not os.path.isfile(file_xml) or not os.path.isfile(file_html):
         return
      if not os.environ.has_key('CCP4'):
         return
      if mode is None:
         return

      jsrview_share = os.path.join(os.environ["CCP4"], "share", "jsrview")

      if mode is not 'i2':
         jsrview_exec = os.path.join(os.environ["CCP4"], "libexec", "jsrview")
         pyrvapi.rvapi_init_document ("ARCIMBOLDO", wrk_path, "ARCIMBOLDO", 1, 7, jsrview_share, None, "index.html", None, None)
      else:
         pyrvapi.rvapi_init_document ("ARCIMBOLDO", wrk_path, "ARCIMBOLDO", 2, 7, jsrview_share, None, None, None, 'i2.xml')
      
      tree = ET.parse(file_xml)
      results = tree.getroot()

      pyrvapi.rvapi_add_header("ARCIMBOLDO")
      pyrvapi.rvapi_add_tab('tab1', 'Report', True)

      self.display_arcimboldo(file_html)
      self.display_data(results)
      self.display_graphs(results)

      if mode is not 'i2':
         self.create_table(results)
         pyrvapi.rvapi_add_tab('tab2', 'Log', True)
         pyrvapi.rvapi_add_section('log', 'Arcimboldo log', 'tab2', 0, 0, 1, 1, True)
         self.display_log(wrk_path)

      pyrvapi.rvapi_flush()

   def search_molecular(self, log):
      return re.search(r'molecular_weight = (.*)', log).group(1)

   def display_log(self, wrk_path):
      file_log = os.path.join(wrk_path,'log')+'.txt'
      file = ''
      text = ''
      file = open(file_log, 'r')
      text = file.read()
      file.close()
      pyrvapi.rvapi_add_text('<pre>'+text+'</pre>', 'log', 0, 0, 1, 1)

   def display_arcimboldo(self,file_html):
      pyrvapi.rvapi_add_section('arcimboldo', 'Arcimboldo', 'tab1', 0, 0, 1, 1, True)

      html_text = '<a href="'+file_html+'"> Check the HTML generated by Arcimboldo</a>'
      pyrvapi.rvapi_add_text(html_text, 'arcimboldo', 0, 0, 1, 1)
      pyrvapi.rvapi_add_text(' ', 'arcimboldo', 1, 0, 1, 1)
      pyrvapi.rvapi_add_text('Please cite: ', 'arcimboldo', 2, 0, 1, 1)
      pyrvapi.rvapi_add_text('Sammito, M., Millan, C., Rodriguez, D. D., M. de Ilarduya, I., Meindl, K., De Marino, I., Petrillo, G., Buey, R. M., de Pereda, J. M., Zeth, K., Sheldrick, G. M. & Uson, I. (2013) Exploiting tertiary structure through local folds for crystallographic phasing. Nature Methods. 10, 1099-1101.', 'arcimboldo', 3, 0, 1, 1)

   def display_data(self, results):
      pyrvapi.rvapi_add_section('data', 'Summary of your data', 'tab1', 0, 0, 1, 1, True)

      if results.find('data/spacegroup') is not None:
          pyrvapi.rvapi_add_text('Spacegroup: '+results.find('data/spacegroup').text, 'data', 0, 0, 1, 1)

      if results.find('data/cell_dim') is not None:
         cell_dim = 'Cell dimensions:'
         if results.find('data/cell_dim/A') is not None:
            cell_dim += ' '+results.find('data/cell_dim/A').text
         if results.find('data/cell_dim/B') is not None:
            cell_dim += ' '+results.find('data/cell_dim/B').text
         if results.find('data/cell_dim/C') is not None:
            cell_dim += ' '+results.find('data/cell_dim/C').text
         if results.find('data/cell_dim/alpha') is not None:
            cell_dim += ' '+results.find('data/cell_dim/alpha').text
         if results.find('data/cell_dim/beta') is not None:
            cell_dim += ' '+results.find('data/cell_dim/beta').text
         if results.find('data/cell_dim/gamma') is not None:
            cell_dim += ' '+results.find('data/cell_dim/gamma').text
         pyrvapi.rvapi_add_text(cell_dim, 'data', 1, 0, 1, 1)

      if results.find('data/resolution') is not None:
         pyrvapi.rvapi_add_text('Resolution: '+results.find('data/resolution').text, 'data', 2, 0, 1, 1)

      if results.find('data/unique_refl') is not None:
         pyrvapi.rvapi_add_text('Number of unique reflections: '+results.find('data/unique_refl').text, 'data', 3, 0, 1, 1)

      pyrvapi.rvapi_flush()

   def display_graphs(self, results):

      pyrvapi.rvapi_add_section('graphs', 'Graphs', 'tab1', 0, 0, 1, 1, True)
      self.create_graph(results)
      pyrvapi.rvapi_flush()

   def create_graph(self, results):
      pyrvapi.rvapi_add_graph("graph", 'graphs', 1, 0, 1, 1)
      pyrvapi.rvapi_add_graph_data ( "data", "graph", "CC")
      pyrvapi.rvapi_add_graph_plot ( "plot", "graph", "Percentage of scattering", "Percentage of scattering", "CC")
      pyrvapi.rvapi_set_plot_legend ( "plot", "graph", "", "")
      x = 'CC'
      y = 'CC percentage of scattering'
      pyrvapi.rvapi_add_plot_line ("plot", "data", "graph", x, y)  
      pyrvapi.rvapi_set_line_options(y, "plot", "data", "graph", "green", "off", "", 2.5, True)
      pyrvapi.rvapi_add_graph_dataset (x, "data", "graph", x, x)
      pyrvapi.rvapi_add_graph_dataset (y,"data","graph", y, y)
      if results.find("backtracing/ALLCC") is not None and results.find("backtracing/ALLCC") is not None: 
         mw = self.search_molecular(results.find("configuration/bor_text").text)
         allnres = json.loads(results.find("backtracing/ALLNRES").text) 
         allcc = json.loads(results.find("backtracing/ALLCC").text)
         if(len(allnres) == len(allcc)):
            for i in range (0, len(allnres)):
               pyrvapi.rvapi_add_graph_real (x, "data", "graph", (int(allnres[i]) * 70 / float(mw)) * 100, "%g")
               pyrvapi.rvapi_add_graph_real (y, "data", "graph", allcc[i], "%g") 

   def fill_table(self, path, tableid, row, col, start, end):
      for i in range (start, end):
         pyrvapi.rvapi_put_table_string(tableid , path[i].text , row , col);
         pyrvapi.rvapi_shape_table_cell(tableid, row, col, "", "font-weight: bold;", "", 1, 1)
         col += 1

   def add_table(self, tableid, fragmentid):
      pyrvapi.rvapi_add_table(tableid, fragmentid, "table", 1, 0, 1, 1, True)
      position = 0
      for i in range(len(self.headers)):
         subheaders_data = (self.headers[i]["subheaders_table"]).split(",")
         pyrvapi.rvapi_put_horz_theader(tableid, self.headers[i]["header_table"], "", i)
         pyrvapi.rvapi_shape_horz_theader ( tableid, i, "", "", 2, len(subheaders_data))
         for j in range(len(subheaders_data)):
            pyrvapi.rvapi_put_table_string(tableid, subheaders_data[j], 1, position+j)
            pyrvapi.rvapi_shape_table_cell(tableid, 1, position+j, "", "", "table-blue-vh", 2, 1)
         position += len(subheaders_data)

   def create_table(self, results):
      pyrvapi.rvapi_add_section('table', 'Tables', 'tab1', 0, 0, 1, 1, True)

      fran = 1
      while 1:
         ensfra = results.find('ens1_frag' + str(fran))
         if ensfra is None:
            break

         tableid = "table" + str(fran)
         fragmentid = "Fragment " + str(fran)            
         self.add_table(tableid, fragmentid)

         if results.find("ens1_frag" + str(fran) + "/FRF") is not None:

            clustn = json.loads(results.find("ens1_frag" + str(fran) + "/FRF/allclus").text)
            numfrag = 3
            for key in clustn:
               keyd = str(key).replace(" ", "").replace("(", "-").replace(")", "-").replace(",", "_").replace("[","-").replace("]", "-")
               path_key = results.find("ens1_frag" + str(fran) + "/FRF/C" + keyd)
               cluster = path_key[0].text
               self.fill_table(path_key, tableid, numfrag, 0, 0, 6)
               if results.find("ens1_frag" + str(fran) + "/FTF") is not None:
                  clustftf = json.loads(results.find("ens1_frag" + str(fran) + "/FTF/allclus").text)
                  if key in clustftf:
                     keyd = str(key).replace(" ", "").replace("(", "-").replace(")", "-").replace(",", "_").replace("[","-").replace("]", "-")
                     path_key = results.find("ens1_frag" + str(fran) + "/FTF/C" + keyd)
                     self.fill_table(path_key, tableid, numfrag, 6, 1, 6)
                     if results.find("ens1_frag" + str(fran) + "/PACK") is not None:
                        clustpack = json.loads(results.find("ens1_frag" + str(fran) + "/PACK/allclus").text)
                        if key in clustpack:
                           keyd = str(key).replace(" ", "").replace("(", "-").replace(")", "-").replace(",", "_").replace("[","-").replace("]", "-")
                           path_key = results.find("ens1_frag" + str(fran) + "/PACK/C" + keyd)
                           self.fill_table(path_key, tableid, numfrag, 11, 1, 6)
                           if results.find("ens1_frag" + str(fran) + "/RNP") is not None:
                              clustrnp = json.loads(results.find("ens1_frag" + str(fran) + "/RNP/allclus").text)
                              if key in clustrnp:
                                 keyd = str(key).replace(" ", "").replace("(", "-").replace(")", "-").replace(",", "_").replace("[","-").replace("]", "-")
                                 path_key = results.find("ens1_frag" + str(fran) + "/RNP/C" + keyd)
                                 self.fill_table(path_key, tableid, numfrag, 16, 1, 4)
                                 if results.find("ens1_frag" + str(fran) + "/INITCC") is not None:
                                    clustinitcc = json.loads(results.find("ens1_frag" + str(fran) + "/INITCC/allclus").text)
                                    if key in clustinitcc:
                                       keyd = str(key).replace(" ", "").replace("(", "-").replace(")", "-").replace(",", "_").replace("[","-").replace("]", "-")
                                       path_key = results.find("ens1_frag" + str(fran) + "/INITCC/C" + keyd)
                                       self.fill_table(path_key, tableid, numfrag, 19, 1, 2)
                                       if results.find("ens1_frag" + str(fran) + "/EXP") is not None:
                                          clusterBest = (results.find("ens1_frag" + str(fran) + "/EXP/Cluster")).text
                                          if(cluster == clusterBest):
                                             path_key = results.find("ens1_frag" + str(fran) + "/EXP")
                                             self.fill_table(path_key, tableid, numfrag, 20, 4, 5)
                                             self.fill_table(path_key, tableid, numfrag, 21, 1, 3)
               numfrag += 1

         fran += 1

      pyrvapi.rvapi_flush()

      return 

class arcimboldo(CPluginScript):

   TASKTITLE = 'Arcimboldo'
   TASKNAME = 'arcimboldo'
   TASKCOMMAND = 'ARCIMBOLDO_LITE'
   TASKVERSION= 0.1
   WHATNEXT = [ 'prosmart_refmac' ]
   MAINTAINER = 'jtvcri@ibmb.csic.es'
   PURGESEARCHLIST = [ 'hklin*.mtz',0 ]

   def genHKL(self, hklin):
      binf = os.path.normpath(os.path.join( CCP4Utils.getCCP4Dir().__str__(), 'bin', 'mtz2hkl' ))
      arglist = ['-f', hklin.__str__()]
      pid = CCP4Modules.PROCESSMANAGER().startProcess(binf, arglist)
      return CCP4Modules.PROCESSMANAGER().getJobData(pid, 'exitCode')


   def generateBor (self, hklin, columns):
      inputData = self.container.inputData
      advancedData = self.container.advancedData
      controlParameters = self.container.controlParameters
      guiAdmin = self.container.guiAdmin

      f_bor = open(os.path.join(self.getWorkDirectory(),'setup.bor'),'w')
      f_bor.write('[CONNECTION]\n')
      f_bor.write('distribute_computing = %s\n' % (controlParameters.ARCIMBOLDO_RUN))
      if controlParameters.RUN_MODE != 'multiprocessing':
         config_file = os.path.join(inputData.CONFIG_FILE.relPath.__str__(),inputData.CONFIG_FILE.baseName.__str__())
         f_bor.write('setup_bor_path = %s\n' % (config_file))

      f_bor.write('[GENERAL]\n')
      f_bor.write('working_directory = %s\n' % (self.getWorkDirectory()))
      f_bor.write('mtz_path = %s\n' % (hklin))
      base = os.path.splitext(hklin)[0]
      f_bor.write('hkl_path = %s.hkl\n' % (base))

      if controlParameters.ARCIMBOLDO_OPTIONS == 'LITE':
         f_bor.write('[ARCIMBOLDO]\n')
         f_bor.write('coiled_coil = %s\n' % (controlParameters.COIL_COILED))
         f_bor.write('rmsd = %6.2f\n' % (controlParameters.LITE_RMSD))
         if controlParameters.LITE_MODELS == 'HELIX':
            f_bor.write('fragment_to_search = %d\n' % (controlParameters.N_FRAGMENTS)) 
            f_bor.write('helix_length = %d\n' % (controlParameters.HELIX_LENGTH)) 

         elif controlParameters.LITE_MODELS == 'CUSTOM':
            f_bor.write('fragment_to_search = %d\n' % (controlParameters.N_FRAGMENTS)) 
            f_bor.write('model_file = %d\n' % (inputData.PDB_LITE)) 

         elif controlParameters.LITE_MODELS == 'HELICES':
            frag_count = 0
            i = 0
            f_bor.write('fragment_to_search = %d\n' % (len(inputData.LITE_HELICES_LIST))) 
            while i < len(inputData.LITE_HELICES_LIST):
               f_bor.write('helix_length_%d = %d\n' % (i+1,inputData.LITE_HELICES_LIST[i])) 
               i += 1

         elif controlParameters.LITE_MODELS == 'CUSTOMS':
            frag_count = 0
            i = 0
            f_bor.write('fragment_to_search = %d\n' % (len(inputData.LITE_CUSTOMS_LIST))) 
            while i < len(inputData.LITE_CUSTOMS_LIST):
               f_bor.write('model_file_%d = %s\n' % (i+1,inputData.LITE_CUSTOMS_LIST[i])) 
               i += 1
     
      elif controlParameters.ARCIMBOLDO_OPTIONS == 'BORGES':
         f_bor.write('[ARCIMBOLDO-BORGES]\n')
         if controlParameters.BORGES_LIBRARY == 'CUSTOM':
            f_bor.write('library_path = %s\n' % (inputData.BORGES_CUSTOM))
         else:
            ccp4_master_home = os.environ.get ( 'CCP4_MASTER', 'not_set' )
            lib_path = os.path.join(ccp4_master_home,'BORGES_LIBS',controlParameters.BORGES_LIBRARY.__str__())
            f_bor.write('library_path = %s\n' % (lib_path))

         if controlParameters.BORGES_GYRE:
            f_bor.write('rotation_model_refinement = %s\n' % (controlParameters.BORGES_GYRE_T))

         if controlParameters.BORGES_GIMBLE:
            f_bor.write('gimble = %s\n' % (controlParameters.BORGES_GIMBLE_T))
                    
      else:
         f_bor.write('[ARCIMBOLDO-SHREDDER]\n')
         f_bor.write('model_file = %s\n' % (inputData.PDB_SHREDDER))
         f_bor.write('rmsd_shedder = %6.2f\n' % (controlParameters.SHREDDER_RMSD))
         if controlParameters.SHREDDER_CONVERT:
            f_bor.write('trim_to_polyala = %s\n' % (controlParameters.SHREDDER_CONVERT_T))
         if controlParameters.SHREDDER_MAKE:
            f_bor.write('bfacnorm = %s\n' % (controlParameters.SHREDDER_MAKE_T))

         f_bor.write('shred_method = %s\n' % (controlParameters.SHREDDER_OPTIONS))
         if controlParameters.SHREDDER_OPTIONS == 'spherical':
            if(controlParameters.SHREDDER_COIL):
               f_bor.write('sphere_definition = default 1 %s 7 4 0.45 0.2\n' % (controlParameters.SHREDDER_COIL_T))
            if(controlParameters.SHREDDER_GYRE):
               f_bor.write('rotation_model_refinement = %s\n' % (controlParameters.SHREDDER_GYRE_T))
            if(controlParameters.SHREDDER_GIMBLE):
               f_bor.write('gimble = %s\n' % (controlParameters.SHREDDER_GIMBLE_T))
            if(controlParameters.SHREDDER_LLG):
               f_bor.write('occ = %s\n' % (controlParameters.SHREDDER_LLG_T))

      #Common configuration values
      f_bor.write('name_job = %s\n' % (guiAdmin.jobTitle))
      f_bor.write('molecular_weight = %10.2f\n' % (controlParameters.MOLECULAR_WEIGHT))
      f_bor.write('number_of_component = %d\n' % (controlParameters.N_COMPONENTS))

      if(controlParameters.TNCS):
         f_bor.write('TNCS = %s\n' % (controlParameters.TNCS))

      if advancedData.SHELXE_LINE.isSet():
         f_bor.write('shelxe_line = %s\n' % (advancedData.SHELXE_LINE))

      if advancedData.KEYWORDS.isSet():
         f_bor.write('%s\n' % (advancedData.KEYWORDS))
        
      if inputData.F_SIGF.contentFlag == 2 or inputData.F_SIGF.contentFlag == 4:
         f_bor.write('f_label = %s\n' % (columns[0]))
         f_bor.write('sigf_label = %s\n' % (columns[1]))
      elif inputData.F_SIGF.contentFlag == 1 or inputData.F_SIGF.contentFlag == 3:
         f_bor.write('i_label = %s\n' % (columns[0]))
         f_bor.write('sigi_label = %s\n' % (columns[1]))

      f_bor.write('[LOCAL]\n')
      f_bor.write('path_local_phaser = %s/bin/phaser\n' % (ccp4_home))
      f_bor.write('path_local_shelxe = %s/bin/shelxe\n' % (ccp4_home))

      f_bor.close()

   def processInputFiles ( self ):
      list_of_stuff = [ ]
      inputData = self.container.inputData

      if inputData.F_SIGF.contentFlag == 1 or inputData.F_SIGF.contentFlag == 3:
         list_of_stuff.append ( [ 'F_SIGF', CCP4XtalData.CObsDataFile.CONTENT_FLAG_IMEAN ] )
      elif inputData.F_SIGF.contentFlag == 2 or inputData.F_SIGF.contentFlag == 4:
         list_of_stuff.append ( [ 'F_SIGF', CCP4XtalData.CObsDataFile.CONTENT_FLAG_FMEAN ] )
        
      self.hklin, self.columns, error = self.makeHklin0 ( list_of_stuff )
      self.columns = self.columns.split(',')
      exitCode = self.genHKL(self.hklin)
      if exitCode != 0:
         return CPluginScript.FAILED

      self.generateBor(self.hklin, self.columns)

      return CPluginScript.SUCCEEDED

   def makeCommandAndScript(self):
      controlParameters = self.container.controlParameters

      if controlParameters.ARCIMBOLDO_OPTIONS == 'LITE':
         self.TASKCOMMAND = 'ARCIMBOLDO_LITE'
      elif controlParameters.ARCIMBOLDO_OPTIONS == 'BORGES':
         self.TASKCOMMAND = 'ARCIMBOLDO_BORGES'
      else:
         self.TASKCOMMAND = 'ARCIMBOLDO_SHREDDER'

      self.appendCommandLine([os.path.join(self.getWorkDirectory(),'setup.bor')])

      return CPluginScript.SUCCEEDED

   def processOutputFiles(self):
      guiAdmin = self.container.guiAdmin
      outputData = self.container.outputData
      
      ar = arcimboldo_rvapi()
      ar.generate_report(self.getWorkDirectory(),str(guiAdmin.jobTitle),'i2')

      pdbout = os.path.join(self.getWorkDirectory(), "best.pdb")
      if os.path.exists(pdbout):
         outputData.XYZOUT.append(outputData.XYZOUT.makeItem())
         outputData.XYZOUT[-1].setFullPath(pdbout)
         outputData.XYZOUT[-1].annotation = 'Best pdb solution'
      else:
         return CPluginScript.FAILED

      return CPluginScript.SUCCEEDED

#=====================================================================================================
#=================================test suite=========================================================
#=====================================================================================================

import unittest
from core.CCP4Utils import getCCP4I2Dir,getTMP

# unit testing asynchronous processes potential tricky but QProcess has option to wait for finished
 
class test_arcimboldo ( unittest.TestCase ) :
   def setUp(self):
      # make all background jobs wait for completion
      PROCESSMANAGER().setWaitForFinished(10000)

   def tearDown(self):
      PROCESSMANAGER().setWaitForFinished(-1)

   def test_arcimboldo(self):
      import os
      inputData =  CScriptDataContainer(name='test_arcimboldo_test',containerType='inputData',initialise=test_arcimboldo.INPUTDATA)
      outputData =  CScriptDataContainer(name='test_arcimboldo_test',containerType='outputData',initialise=test_arcimboldo.OUTPUTDATA)
      try:
         inputData.importXML(os.path.join(getCCP4I2Dir(),'wrappers','test_arcimboldo','test_data','test_arcimboldo_test_1.def.xml'))
      except CException as e:
         self.fail(e.errorType)
      try:
         outputData.importXML(os.path.join(getCCP4I2Dir(),'wrappers','test_arcimboldo','test_data','test_arcimboldo_test_1.def.xml'))
      except CException as e:
         self.fail(e.errorType)
      
      wrapper = test_arcimboldo()
      pid = wrapper.process()

   def testSuite():
      suite = unittest.TestLoader().loadTestsFromTestCase(test_test_arcimboldo)
      return suite

   def runAllTests():
      suite = testSuite()
      unittest.TextTestRunner(verbosity=2).run(suite)