# PDB file, coot script file handling.
#
# Author       : Chris X. Weichenberger
# Created      : 2012-02-24
# License
# Twilight is a GUI for displaying ligands incompatible with electron density.
#    Copyright (C) 2012  Christian X. Weichenberger, Edwin Pozharski, and
#    Bernhard Rupp.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    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 General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Modifications:
# 2012-03-15 (_X) Sorted list of het groups.
# 2012-07-24 (CW) Import of PDB and EDS data directly from EDS.
# 2012-08-01 (CW) Fix for TNT refinement program parsing of R-values.
# 2012-08-02 (CW) Compute the ligands' mol. weight when parsing the PDB file.
# 2012-10-10 (CW) Better discrimination between failures. Loading from PDB and
#  from EDS in order to find out if a PDB entry is missing due to electron
#  density unavailability.
# 2012-10-15 (CW) Use of deposit date instead of release date.

import os
import sys
import gzip
import urllib
import shutil
import tempfile

from misc     import NOVALIDFLOATSTR, NONLIGANDS, INVALIDPDBDATE;
from ligand   import CLigand;
from eds      import CEDSFile, EDSRESNMIDX, EDSRSCCIDX, EDSOWABIDX;

from files import Debug;

# Error message constants.
# PDB file retrieved, everything ok.
PDB_OK                  = 0;
# Could not establish internet connection.
PDB_NO_CONNECTION       = 1;
# Could not find entry/file on PDB server.
PDB_PDBENTRY_NOT_EXISTS = 2;
# Could not find entry/file on EDS server.
PDB_EDSENTRY_NOT_EXISTS = 3;
# Invalid Mapman file from EDS.
PDB_EDS_INVALID_STATLIS = 4;
# No ligand data present in PDB file.
PDB_NO_LIGANDS_FOUND    = 5;

class CPDBFile(list):
    EDSBaseURL= "http://eds.bmc.uu.se/eds/dfs/%s/%s/";
    PDBBaseURL= "http://www.rcsb.org/pdb/files/%s.pdb";
    def __init__(s):
        "Handle PDB/EDS data. dataDir may point to a PDB archive mirror." 
        s._PDBDir         = "";
        s._EDSDir         = "";
        s._CootDir        = "";
        s._CootExec       = "";
        s._CurrentPDBID   = "";
        s._ShellScriptExt = ".bat" if sys.platform=="win32" else ".sh";
        # this option should be overridden by a call from the config module.
        s._ScriptingHost = "python";
    
    def _GetLocalPDBFileName(s, pdbID):
        return os.path.join(s._CootDir, pdbID + ".ent");
    
    def _GetMirrorPDBFileName(s, pdbID):
        return os.path.join(s._PDBDir, pdbID[1:3], "pdb%s.ent.gz" % pdbID)

    def _GetLocalMTZFileName(s, pdbID):
        return os.path.join(s._CootDir, pdbID + "_sigmaa.mtz");

    def _GetMirrorEDSStatLisFileName(s, pdbID):
        return os.path.join(s._EDSDir, pdbID + "_stat.lis");
    
    def _GetLocalEDSStatLisFileName(s, pdbID):
        return os.path.join(s._CootDir, pdbID + "_stat.lis");
        
    def _GetLocalCootScriptFileName(s, pdbID):
        ext = ".py" if s._ScriptingHost=="python" else ".scm";
        return os.path.join(s._CootDir, pdbID + ext);

    def _GetLocalShellScriptFileName(s, pdbID):
        return os.path.join(s._CootDir, pdbID + s._ShellScriptExt);
    
    def _GroupLines(s):
        s._Records = {};
        for line in s:
            s._Records.setdefault(line[0:6], []).append(line);
    
    def _DetermineHetGroups(s):
        s._HetGroups = {};
        for line in s._Records.get("HETATM", []):
            rnr = line[21:27];
            rnm = line[17:20];
            if( rnm in NONLIGANDS ): continue 
            if( s._HetGroups.get(rnr, None)==None ):
                chemFormula = "";
                for chemLine in s._Records.get("FORMUL", []):
                    if( chemLine[12:15]==rnm ):
                        chemFormula = chemLine[19:70].strip();
                ligand = CLigand();
                ligand.SetResNr(rnr);
                ligand.Set3LetterName(rnm);
                ligand.SetChemFormula(chemFormula);
                ligand.SetMolWt( ligand.ComputeMolWeight() );
                s._HetGroups[rnr] = ligand;
                
    def _ConvertPDBDate(s, date):
        m2m  = {"JAN": "01", "FEB": "02", "MAR": "03", "APR": "04",
                "MAY": "05", "JUN": "06", "JUL": "07", "AUG": "08",
                "SEP": "09", "OCT": "10", "NOV": "11", "DEC": "12"};
        try:
            dd, mm, yy = date.split("-");
        except ValueError:
            # If the date format does not produce three values, take the
            # invalid date.
            dd, mm, yy = INVALIDPDBDATE.split("-");
        yyprefix = "19" if int(yy)>=60 else "20";
        return "%s-%s-%s" % (yyprefix+yy, m2m.get(mm, "00"), dd);
    
    def SetPDBDir(s, pdbDir):
        """Directory name or empty @c String."""
        # XXX this needs some more work...
        s._PDBDir = pdbDir;

    def SetEDSDir(s, edsDir):
        """Directory name or empty @c String."""
        s._EDSDir  = edsDir;
        
    def SetCootTmpDir(s, cootDir):
        s._CootDir = cootDir;
    
    def SetCootExec(s, cootExec):
        s._CootExec = cootExec;
    
    def SetCootScriptingHost(s, lang):
        s._ScriptingHost = lang;
    
    def GetAllHetGroups(s):
        """
        @return  @c List of @c Tuple (residue number, CLigand object).
        """
        res = [(rnr, ligand) for rnr, ligand in s._HetGroups.items() ];
        res.sort();
        return res;
    
    def GetCenterAtom(s, resNr):
        for line in s:
            if( line[:6]!="HETATM" ): continue ;
            if( line[21:27]==resNr ):
                # atom name of first atom of this group.
                return line[12:16]; 
        return "";

    def GetReleaseDate(s):
        date = INVALIDPDBDATE;
        for line in s._Records.get("REVDAT", []):
            if( line[:11]=="REVDAT   1 " ):
                date = line[13:22];
        return s._ConvertPDBDate(date);
    
    def GetDepositDate(s):
        noHeader = ["HEADER    XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX   %s   9XXX              " % INVALIDPDBDATE,]
        date = s._Records.get("HEADER", noHeader)[0][50:59];
        return s._ConvertPDBDate(date);
        
    def GetResolution(s):
        """Get experimental resolution. Also for non-xray stuff."""
        resol   = NOVALIDFLOATSTR;
        for line in s._Records.get("REMARK", []):
            if( line[:11]=="REMARK   2 " ):
                parts = line.split();
                try:
                    resol = parts[3];
                except IndexError:
                    resol = NOVALIDFLOATSTR;
                try:
                    float(resol);
                except ValueError:
                    resol = NOVALIDFLOATSTR;
        return float(resol);

    def GetR(s):
        def convertR(line):
            try:
                val = float(line.split(":")[1]);
            except:
                val = float(NOVALIDFLOATSTR);
            return val;
        # There are multiple blocks addressing the R values. We need to be
        # sure to parse the right one. It starts with a REMARK   3  FIT TO
        # DATA USED IN REFINEMENT. Once we have parsed it, another block
        # (FIT/AGREEMENT for example) may not be parsed.
        rValues       = [float(NOVALIDFLOATSTR), float(NOVALIDFLOATSTR),
                         float(NOVALIDFLOATSTR)];
        haveRefineSec = False;
        fitTok        = "REMARK   3  FIT";
        fitTNTTok     = "REMARK   3  USI";
        refineTok     = "REMARK   3  FIT TO DATA USED IN REFINEMENT";
        refineTNTTok  = "REMARK   3  USING DATA ABOVE SIGMA CUTOFF.";
        rWrkOrObsTok  = "REMARK   3   R VALUE";
        rFreTok       = "REMARK   3   FREE R VALUE     ";
        for line in s._Records.get("REMARK", []):
            if( line[:15] in (fitTok, fitTNTTok) and haveRefineSec ):
                haveRefineSec = False;
                continue ;
            if( line[:42] in (refineTok, refineTNTTok) ):
                haveRefineSec = True;
            if(   haveRefineSec and line[:20]==rWrkOrObsTok ):
                words = line.split();
                if( "(WORKING" in words and "TEST" in words ):
                    rValues[0] = convertR(line);
                elif( "(WORKING" in words and "TEST" not in words ):
                    rValues[1] = convertR(line);
            elif( haveRefineSec and line[:30]==rFreTok ):
                rValues[2] = convertR(line);
        return rValues;
    
    def GetRObs(s):
        return s.GetR()[0];

    def GetRWork(s):
        return s.GetR()[1];

    def GetRFree(s):
        return s.GetR()[2];

    def GetRefinementSoftware(s):
        refProgs = [];
        for line in s._Records.get("REMARK", []):
            if( line[:26]=="REMARK   3   PROGRAM     :" ):
                progStr = line[27:].strip();
                for programPlusVersion in progStr.split(", "):
                    components = programPlusVersion.split(" ");
                    # this gets rid of version numbers.
                    refProgs.append(components[0]);
        return ",".join(refProgs);

    def GetJournal(s):
        for line in s._Records.get("JRNL  ", []):
            if( line[:16]=="JRNL        REF " ):
                if( line[:34]!="JRNL        REF    TO BE PUBLISHED" ):
                    return line[19:47].strip();
        return "To be published";

    def GetPubMedID(s):
        for line in s._Records.get("JRNL  ", []):
            if( line[:16]=="JRNL        PMID" ):
                try:
                    # I saw one (now already updated) PDB file with missing
                    # PMID, even though the line was present (2w1p).
                    return line.split()[2];
                except IndexError:
                    pass ;
        return "-";

    def FetchPDB(s, pdbID):
        """Fetch PDB file from EDS server or from local cache.
        @return  @c PDB_OK if fetched and content is a PDB file.
         PDB_NO_CONNECTION if the internet connection could not be established
         to the server.
         PDB_PDBENTRY_NOT_EXISTS if the PDB entry could not be found.
        """
        if( s._CurrentPDBID==pdbID ):
            # There is nothing to do, the file has already been loaded into
            # this object.
            return PDB_OK;
        # This will be set with to the real PDB ID when loaded successfully.
        s._CurrentPDBID = "";
        mirrorPDBFN     = s._GetMirrorPDBFileName(pdbID);
        localPDBFN      = s._GetLocalPDBFileName(pdbID);
        # First, let's see if we have already fetched that file previously.
        if(   os.path.isfile(localPDBFN) ):
            fh = open(localPDBFN);
        # See if there is a local PDB mirror installation. If yes, it must not
        # be complete, so check if the file is there.
        elif( os.path.isfile(mirrorPDBFN) ):
            fh = gzip.GzipFile(mirrorPDBFN);
        # Ok, neither in the local cache, nor in a PDB mirror, get it online.
        else:
            # For speed reasons, we download the PDB file from RCSB as a gzip
            # compressed file and uncompress it to the final local file.
            middle        = pdbID[1:3];
            pdbURL        = s.PDBBaseURL % pdbID;
            try:
                urllib.urlretrieve(pdbURL, localPDBFN);
            except:
                return PDB_NO_CONNECTION;
            fh  = open(localPDBFN);
        s[:]= [line for line in fh];
        fh.close();
        s._GroupLines();
        # After the grouping call, we have split by the record
        # names. Let's see if we have PDB-compliant records, especially
        # those we are interested in: a) hetgroups for ligands, and
        # standard amino acids such that the hetgroups become ligands in a
        # protein chain.
        if( "ATOM  " not in s._Records and "HETATM" not in s._Records ):
            # Obviously this is not the case. This cannot be a PDB file
            # then. It might be an error message (page) from PDB. Wave
            # goodbye. 
            s[:] = [];
            return PDB_PDBENTRY_NOT_EXISTS;
        s._DetermineHetGroups();
        fh = open(localPDBFN, "w");
        fh.write("".join(s));
        fh.close();
        s._CurrentPDBID = pdbID;
        return PDB_OK;

    def FetchMTZ(s, pdbID):
        """Fetch an MTZ file from the EDS server.
        @return  @c PDB_OK or PDB_NO_CONNECTION.
        """
        middle     = pdbID[1:3];
        mtzURL     = s.EDSBaseURL % (middle, pdbID) + pdbID + "_sigmaa.mtz";
        localMTZFN = s._GetLocalMTZFileName(pdbID);
        if( not os.path.isfile(localMTZFN) ):
            try:
                urllib.urlretrieve(mtzURL, localMTZFN);
            except:
                return PDB_NO_CONNECTION;
        # Get the magic string from the file. Should be "MTZ" if it's a real
        # MTZ file. Otherwise it's an error message from EDS.
        fh = open(localMTZFN, "rb");
        mtz3Bytes = fh.read(3);
        fh.close();
        if( mtz3Bytes!="MTZ" ):
            return PDB_EDSENTRY_NOT_EXISTS;
        return PDB_OK;

    def FetchEDSStatLis(s, pdbID):
        """Fetch an '_stat.lis' file from the EDS server.
        @return  PDB_OK, PDB_NO_CONNECTION, or PDB_EDSENTRY_NOT_EXISTS.
        """
        # If we load a new stat.lis file on top of a different PDB structure,
        # we would mess up things. Therefore, invalidate the current PDB ID,
        # if it is not the same.
        if( s._CurrentPDBID!=pdbID ):
            s._CurrentPDBID = pdbID;
        middle    = pdbID[1:3];
        mirrorSLFN= s._GetMirrorEDSStatLisFileName(pdbID);
        localSLFN = s._GetLocalEDSStatLisFileName(pdbID);
        slURL     = s.EDSBaseURL % (middle, pdbID) + pdbID + "_stat.lis";
        if( not os.path.isfile(localSLFN) ):
            if( not os.path.isfile(mirrorSLFN) ):
                # XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
                #return False
                try:
                    urllib.urlretrieve(slURL, localSLFN);
                except:
                    # This is only a connection error. If the file has not been
                    # found on the EDS server, a small web page is going to be
                    # sent with a 404 Not Found title.
                    return PDB_NO_CONNECTION;
            else:
                shutil.copyfile(mirrorSLFN, localSLFN);
        
        eds = CEDSFile();
        # Check if the file is really containing data and not only an error
        # message from the server.
        return PDB_OK if eds.is_valid(localSLFN) else PDB_EDSENTRY_NOT_EXISTS;

    def _PopulateLigandData(s, pdbID):
        """Combine ligand data from the PDB entry and from EDS.
        Call this method here only, if both the PDB file and the Stat list
        file have been fetched before and nothing else in between.
        @return  True if things went ok, False if the pdbID is not the current
         PDB ID that is stored in this object.
        """
        if( s._CurrentPDBID!=pdbID ):
            return False;
        eds   = CEDSFile();
        slFFN = s._GetLocalEDSStatLisFileName(pdbID);
        eds.parse(pdbFFN=slFFN);
        # Data should be ok, since this is an internal function that relies on
        # sanity check of the _stat.lis file before calling this function.
        for rnr, ligand in s._HetGroups.items():
            edsData = eds.get(rnr, None);
            if( edsData!=None ):
                ligand.SetRSCC(edsData[EDSRSCCIDX]);
                ligand.SetOWAB(edsData[EDSOWABIDX]);
        return True;
        
    def FetchPDBnStatLis(s, pdbID):
        """Utility function to get all data at once.
        @return  PDB_OK if everything went ok, files have been loaded from PDB
         and EDS. Otherwise, an error code as defined at the beginning of this
         file will be returned.
        """
        succ = s.FetchPDB(pdbID);
        if( succ!=PDB_OK ):
            return succ;
        succ = s.FetchEDSStatLis(pdbID);
        if( succ!=PDB_OK ):
            return succ;
        succ = s._PopulateLigandData(pdbID);
        if( not succ ):
            return PDB_EDS_INVALID_STATLIS;
        if( len(s._HetGroups)==0 ):
            return PDB_NO_LIGANDS_FOUND;
        return 0;
        
    def WriteCootScript(s, pdbID, resNr, width, height, sigma2FoFc, sigmaFoFc,
                                                        reslist=False):
        # Initially, I started to write guile scripts, but Win$ coot does not
        # have guile, only Python. Converged to Python scripting then, runs on
        # both Linux and Win$.
        # (make-and-draw-map "%s" "FOFCWT" "PHFOFCWT" "" 0 1)
        scriptTemplScm='''(read-pdb "%s")
(auto-read-make-and-draw-maps "%s")
(set-go-to-atom-chain-residue-atom-name "%s" %s "%s")
(set-clipping-front -10.00)
(set-clipping-back -10.00)
(set-graphics-window-size %d %d)
(set-contour-level-in-sigma 1 %3.2f)
(set-contour-level-in-sigma 2 %3.2f)
''';
        pyNavScript = '''
def go2next():
    global cblig
    global badlignum
    cblig = cblig+1 if cblig<badlignum-1 else 0
    set_go_to_atom_chain_residue_atom_name(blchid[cblig], 
                                           blresn[cblig], 
                                           blatom[cblig])
def go2prev():
    global cblig
    global badlignum
    cblig = cblig-1 if cblig>1 else badlignum-1
    set_go_to_atom_chain_residue_atom_name(blchid[cblig], 
                                           blresn[cblig], 
                                           blatom[cblig])
add_key_binding("Next bad ligand", "g", lambda: go2next())
add_key_binding("Previous bad ligand", "v", lambda: go2prev())
cblig=0
blchid, blresn, blatom = [], [], []
'''
        # we need to use raw strings, as the path separator on windows is the
        # escape character...
        scriptTemplPy='''read_pdb(r"%s")
auto_read_make_and_draw_maps(r"%s")
set_go_to_atom_chain_residue_atom_name("%s", %s, "%s")
set_clipping_front(-10.00)
set_clipping_back(-10.00)
set_graphics_window_size(%d, %d)
set_contour_level_in_sigma(1, %3.2f)
set_contour_level_in_sigma(2, %3.2f)
''';

        scriptTempl = scriptTemplPy if s._ScriptingHost=="python"\
                      else scriptTemplScm;
        chnNm  = resNr[0];
        rnr    = resNr[1:5];
        mtzFFN = s._GetLocalMTZFileName(pdbID);
        pdbFFN = s._GetLocalPDBFileName(pdbID);
        scmFFN = s._GetLocalCootScriptFileName(pdbID);
        atomNm = s.GetCenterAtom(resNr);
        script = scriptTempl % (pdbFFN, mtzFFN, chnNm, rnr, atomNm,
                                width, height, sigma2FoFc, sigmaFoFc);
        if reslist and s._ScriptingHost=="python":
            for residue in reslist:
                pyNavScript += 'blchid.append("%s")\n' % residue[0]
                pyNavScript += 'blresn.append(%s)\n' % residue[1:5]
                pyNavScript += 'blatom.append("%s")\n' % s.GetCenterAtom(residue)
            pyNavScript += 'badlignum = len(blchid)'
            script += pyNavScript
        fh = open(scmFFN, "w");
        fh.write(script);
        fh.close();
        return scmFFN;

    def WriteShellScript(s, pdbID):
        cootTmpFiles = "0-coot.state.scm 0-coot.state.py 0-coot-history.scm 0-coot-history.py";
        if( sys.platform=="win32" ):
            # win$ has .bat files and we need to quote everyting due to spaces
            # in path names...
            # do not know why, but we need to delete the temp files first and
            # then start coot. The other way round did not work. *scratch*
            scriptTempl  = "del /Q /F %s" + os.linesep;
            scriptTempl += '"' + s._CootExec + '"' + ' -s "%s"' + os.linesep;
        else:
            # some Unix flavor
            scriptTempl = '#/bin/sh' + os.linesep;
            scriptTempl += 'rm -f %s' + os.linesep;
            scriptTempl += s._CootExec + \
                           ' -s "%s" >/dev/null 2>/dev/null'+os.linesep;
        shFFN  = s._GetLocalShellScriptFileName(pdbID);
        scmFFN = s._GetLocalCootScriptFileName(pdbID);
        script = scriptTempl % (cootTmpFiles, scmFFN);
        #Debug(shFFN);
        fh     = open(shFFN, "w");
        fh.write(script);
        fh.close();
        os.chmod(shFFN, 0755);
        return shFFN;
        


if( __name__=="__main__" ):
    # Large scale testing. Checked all ~50,000 PDB entries with an associated
    # EDS entry.
    from misc import CalcCombinedScore;
    obj = CPDBFile();
    obj.SetPDBDir("/Users/chris/Data/PDB");
    obj.SetEDSDir("./data/EDSDB");
    obj.SetCootTmpDir("/Users/chris/tmp/coot-data");
    print "PDB ID\tLigNm\tRNR\tRSCC\tOWAB\tMolWt\tChrg\tNrHeavy\tNrHyd\tNrAtms\tReleaseDate\tResol\tRwrk\tRFree\tRobs\tSftwre\tJrnl\tPMID\tClass\tScore"
    startFrom = "1n08"; #None;
    for pdbID in open("/tmp/pdbXraylist.tsv"):
        pdbID = pdbID.strip();
        if( startFrom!=None and pdbID!=startFrom ):
            continue ;
        startFrom = None;
        succ=obj.FetchPDBnStatLis(pdbID);
        if( succ==0 ):
            hetGroups = obj.GetAllHetGroups();
            dpl = [rnr for rnr, ligand in hetGroups];
            obj.WriteCootScript(pdbID, rnr, 500, 450, 1.0, 3.0, reslist=dpl);
            obj.WriteShellScript(pdbID)
            for rnr, ligand in hetGroups:
                score = CalcCombinedScore(ligand.GetRSCC(),
                                          obj.GetResolution());
                ligand.SetCombinedScore(score);
                row = [pdbID,
                       ligand.Get3LetterName(),
                       rnr,
                       "%4.3f" % ligand.GetRSCC(),
                       "%4.2f" % ligand.GetOWAB(),
                       "%5.1f" % ligand.GetMolWt(),
                       "%d" % ligand.GetCharge(),
                       "%d" % ligand.GetNrHeavyAtoms(),
                       "%d" % ligand.GetNrHydrogenAtoms(),
                       "%d" % ligand.GetNrAtoms(),
                       obj.GetDepositDate(),
                       "%3.2f" % obj.GetResolution(), 
                       "%4.3f" % obj.GetRWork(), 
                       "%4.3f" % obj.GetRFree(), 
                       "%4.3f" % obj.GetRObs(),
                       obj.GetRefinementSoftware(),
                       obj.GetJournal(),
                       obj.GetPubMedID(),
                       ligand.GetClassification(),
                       "%4.3f" % ligand.GetCombinedScore()];
                print "\t".join(row);
                   
        else:
            if( succ in (1,2,4) ):
                sys.stderr.write("Could not fetch: %s\n" % pdbID);
