# IO of ligand table.
#
# 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 (CW) Sorted list of het groups, Ed's get for all bad valid ligands.
# 2012-06-14 (CW) Removed truncation of Jrnl column to have full journal name.
# 2012-07-27 (CW) Adding of ligands.
# 2012-10-15 (CW) Use of deposit date instead of release date.
# 2012-10-14 (CW) Rewrote meaning of bad ligands. Good and bad ones.
# 2012-10-16 (CW) Complex sorting for an arbitrary column.
#

import csv

from misc import CalcCombinedScore, VALID_BAD, VALID_GOOD, VALID_ADDED,\
     VALID_NEITHER, RSCCBADTH, RSCCGOODTH;
from files import Output;

# Number of columns in this file.
NR_COLUMNS = 22;

class CEDSTableData(list):
    # Here we have some class variables.
    PDBIDIDX, LIGNMIDX, RESNRIDX, RSCCIDX, OWABIDX, MOLWTIDX, CHARGEIDX, \
    HEAVYIDX, HYDROGIDX, ATOMSIDX, DATIDX, RELIDX, RWORKIDX, RFREEIDX, \
    ROBSIDX, SWIDX, JRNLIDX, PMIDIDX, CLASSIDX, VALIDIDX, SCOREIDX, \
    COMMENTIDX =\
    range(NR_COLUMNS);
    ColumnTypes = (str   , str , str , float , float , float , int   , 
                   int   , int , int , str   , float , float , float ,
                   float , str , str , str   , str   , str   , float ,
                   str    );

    def __init__(s):
        s._Header     = ();
        # Default sorting: by combined score and descending.
        s.SetSortingScheme(s.SCOREIDX, False);
        # Default ligands on display: bad
        s.SetFocusOnBad()
    
    def _LigandInFocus(s, rscc):
        if( s._FocusIsBad ):
            return rscc<RSCCBADTH;
        else:
            return rscc>=RSCCGOODTH;
            
    def IsValidTableFile(s, filename):
        """Checks if the file looks like an EDS table."""
        fh     = open(filename);
        header = fh.readline();
        fh.close();
        if( header[:5]=="PDBID" ):
            # looks good, let's check the number of columns...
            header = header.split("\t");
            if( len(header)==NR_COLUMNS ):
                return True;
        return False;
        
    def ReadTable(s, filename):
        """Opens and reads an EDS table.
        @param filename  Full path file name.
        """
        Output("Reading table ``%s''..." % filename);
        fh     = open(filename);
        reader = csv.reader(fh, delimiter="\t");
        s._Header = reader.next();
        s._All = [l for l in reader];
        fh.close();
        s._ComplexSort();
        s._ReIndex();

    def SetFocusOnBad(s):
        """Gather bad ligands. Call this prior to loading a file!"""
        s._FocusIsBad = True;
    
    def SetFocusOnGood(s):
        """Gather good ligands. Call this prior to loading a file!"""
        s._FocusIsBad = False;
    
    def ToggleFocus(s):
        """Utility function to invert the outcome of the Get*Ligands methods.
        """
        s._FocusIsBad = not s._FocusIsBad;
    
    def SetSortingScheme(s, sortCol, ascending):
        """Set the sorting column and order.
        @param sortCol  Column index.
        @param ascending @c True or @c False.
        @return @c False if sortCol was incorrect, @c True, if ok.
        """
        if( sortCol<0 or sortCol>=s.COMMENTIDX ):
            return False;
        s._SortColIdx = sortCol;
        s._SortAsc    = ascending;
        
    def IsFocusOnBad(s):
        """What type of ligands do we display?
        @return @c True, bad ones. @c False, good ones.
        """
        return s._FocusIsBad;
    
    def IsManuallyImported(s, pos):
        return s[pos][s.VALIDIDX]==VALID_ADDED;
    
    def GetHeader(s):
        return s._Header;

    def GetPDBID(s, pos):
        return s[pos][s.PDBIDIDX];
    
    def GetResNr(s, pos):
        return s[pos][s.RESNRIDX];

    def GetLigNm(s, pos):
        return s[pos][s.LIGNMIDX];
    
    # formerly: GetAllValidLigands
    def GetAllDisplayLigands(s, pos):
        """
        Get all ligands which should be displayed in the table for this PDB
        entry.
        """ 
        pdbID = s.GetPDBID(pos);
        wanted= [VALID_ADDED, ];
        wanted.append(VALID_BAD if s._FocusIsBad else VALID_GOOD);
        res   = [item[s.RESNRIDX] for item in s._All
                 if item[s.PDBIDIDX]==pdbID and item[s.VALIDIDX] in wanted];
        res.sort();
        return res;

    # formerly: GetAllBadRSCCLigands
    def GetAllRSCCFocusLigands(s, pos):
        """Get all ligands which are currently ligands of desired quality.
        This depends on whether we  want to display ligands with bad RSCC
        values or with good RSCC values.
        """
        pdbID = s.GetPDBID(pos);
        res   = [item[s.RESNRIDX] for item in s._All
                 if item[s.PDBIDIDX]==pdbID and
                 s._LigandInFocus(float(item[s.RSCCIDX]))];
        res.sort();
        return res;
    
    # formerly: GetAllValidBadRSCCLigands
    def GetAllDisplayRSCCFocusLigands(s, pos, fsorted=True):
        """Get all ligands that should be displayed in the table and meet the
        desired RSCC criterion (either good or bad values, depending of the
        focus).
        """
        pdbID  = s.GetPDBID(pos);
        wanted = [VALID_ADDED, ];
        wanted.append(VALID_BAD if s._FocusIsBad else VALID_GOOD);
        res   = [item[s.RESNRIDX] for item in s._All
                 if item[s.PDBIDIDX]==pdbID and
                 s._LigandInFocus(float(item[s.RSCCIDX])) and
                 item[s.VALIDIDX] in wanted];
        if fsorted:
            res.sort();
        return res;
        
    def GetPreviousInPDB(s, pos):
        pdbID = s.GetPDBID(pos);
        for i in range(pos-1,-1,-1):
            if( s[i][s.PDBIDIDX]==pdbID ):
                return i;
        return pos;

    def ConvertRow(s, i):
        """Convert the text in each row to its natural data types."""
        return [dType(col) for dType, col in zip(s.ColumnTypes, s[i])];
        
    def WriteTSV(s, filename, verbose=False):
        """Save file to given file name.
        @retun  @c True, if everything went ok. @c False if an error during
        file opening occurred.
        """
        if( verbose ):
            Output("\rWriting table ``%s''..." % filename);
        for idx, entry in zip(s._Idx, s):
            s._All[idx][s.COMMENTIDX] = entry[s.COMMENTIDX];
        try:
            fh      = open(filename, "w");
            content = "\n".join(["\t".join(line)
                                 for line in [s.GetHeader(),]+s._All]);
            fh.write(content);
            fh.close();
            return True;
        except IOError:
            return False;

    def _ComplexSort(s):
        """Sorts s._All in a meaningful way.
        That is, PDB entries are grouped. Within each PDB group, those rows
        with a valid column are sorted by Score and ranked top. After them,
        columns with valid code VALID_NEITHER, sorted by Score, are added to
        the group. This guarantees that the highest Score is on top of the whole
        group. All groups are again sorted by the Score of its first row,
        therefore ranking the worst ligands per group (ie. PDB entry) on top.
        """
        desc     = lambda x,y: -(x>y);
        asc      = lambda x,y: -(x<y);
        sortFunc = asc if s._SortAsc else desc;
        # sorting within a group, this is by Score.
        perGroup = lambda x,y,i=s._SortColIdx: sortFunc(float(x[i]),float(y[i]));
        # sorting of the groups, this is by the Score of the first row.
        wholeTab = lambda x,y,i=s._SortColIdx: sortFunc(float(x[1][0][i]), 
                                        float(y[1][0][i]));
        # Partition data set into a valid and bad VALID_BAD, a valid and good
        # VALID_GOOD, an added VALID_ADDED, and not valid (ie. neither good
        # nor bad, unexciting) VALID_NEITHER column set.
        pdbBuckets = {VALID_ADDED: {}, VALID_BAD: {}, VALID_GOOD: {},
                      VALID_NEITHER: {}};
        for row in s._All:
            validCode = row[s.VALIDIDX];
            pdbBuckets[validCode].setdefault(row[s.PDBIDIDX], []).append(row[:]);
        # sort the two data sets by Score.
        for validCode, target in pdbBuckets.items():
            for rows in target.values():
                rows.sort(perGroup);
        
        # Merge data sets again.
        # On the very top we put the manually added PDB entries.
        # Below, we place those groups with the highest Score for the valid
        # columns (per group!). Append not valid rows after all valid rows.
        inFocus, outFocus = [VALID_BAD, VALID_GOOD] if s._FocusIsBad else \
                            [VALID_GOOD, VALID_BAD];
        allItems = pdbBuckets[inFocus].items();
        allItems.sort(wholeTab);
        newList = [];
        for rows in pdbBuckets[VALID_ADDED].values():
            newList.extend(rows);
        seen    = set();
        for key, rows in allItems:
            # first the valid rows in focus, worst on top of this group.
            newList.extend(rows);
            # then add those which are currently not in the focus
            newList.extend(pdbBuckets[outFocus].get(key, []));
            # and then append any not valid rows.
            newList.extend(pdbBuckets[VALID_NEITHER].get(key, []));
            seen.add(key);
        # finally complete the list with all groups that did not have a valid
        # row at all.
        for key, rows in pdbBuckets[VALID_NEITHER].items():
            if( key not in seen ):
                newList.extend(rows);
        # replace the original table with this sorted table.
        s._All = newList;

    def _ReIndex(s):
        """Create internal indices.
        This is needed when data have been loaded or freshly added. 
        """
        wanted= [VALID_ADDED,];
        wanted.append(VALID_BAD if s._FocusIsBad else VALID_GOOD);
        s._Idx = [n for n,l in enumerate(s._All) if l[s.VALIDIDX] in wanted];
        s[:] = [l[:] for l in s._All if l[s.VALIDIDX] in wanted];
        
    def AddLigand(s, pdbID, pdbFileObj, ligandObj):
        """Adds a ligand to the top of this list.

        @param pdbID  PDB four letter code.
        @param pdbFileObj  @c CPDBFile object.
        @param ligandObj  @c CLigand object.
        """
        lig = ligandObj;
        pdb = pdbFileObj;
        score = "%4.3f" % CalcCombinedScore(lig.GetRSCC(),pdb.GetResolution());
        elem = [pdbID,
                lig.Get3LetterName(),
                lig.GetResNr(),
                "%3.3f" % lig.GetRSCC(),
                "%4.2f" % lig.GetOWAB(),
                "%5.1f" % lig.GetMolWt(),
                "%2d" % lig.GetCharge(),
                "%3d" % lig.GetNrHeavyAtoms(),
                "%3d" % lig.GetNrHydrogenAtoms(),
                "%3d" % lig.GetNrAtoms(),
                pdb.GetDepositDate(),
                "%3.2f" % pdb.GetResolution(),
                "%4.3f" % pdb.GetRWork(),
                "%4.3f" % pdb.GetRFree(),
                "%4.3f" % pdb.GetRObs(),
                pdb.GetRefinementSoftware(),
                pdb.GetJournal(),
                pdb.GetPubMedID(),
                lig.GetClassification(),
                VALID_ADDED,
                score,
                ""];
        # Re-indexing is guearanteed to work when ligand data has been entered
        # at the top of the list.
        s._All.insert(0, elem);
        s._ReIndex();

    def RemoveLigand(s, pdbID, resNr):
        # probably we need this with an index instead of pdb/rnr...
        """Remove a ligand from the List and the file.
        @param pdbID  PDB four letter code.
        @param resNr  Residue number with chain name and insertion code.
        """
        for i, l in enumerate(s._All):
            if( l[s.PDBIDIDX]==pdbID and l[s.RESNRIDX]==resNr ):
                del s._All[i];
                break ;
        s._ReIndex();

if( __name__=="__main__" ):
    obj=CEDSTableData();
    obj.ReadTable("ligands-2012-01-15.tsv");
    for line in obj[:20]:
        print line;
