#!/usr/bin/env python2.6
# -*- coding: utf-8 -*-

import os, sys

from p3d.protein import Protein

from operator import attrgetter

from nomenclature import metalList
from metalSite import MetalSite

#===============================================================================
"""CLASS"""
#===============================================================================

class SourcePDB(object):

    def __init__(self, pathPdb, maxDist=5.0):
        
        self.maxDist = maxDist
        
        self.code = pathPdb.split('/')[-1].split('.')[0][0:4].lower()
        try:
            self.pdb = Protein(pathPdb)
        except:
            raise Exception('Unable to open %s' % self.code)
            sys.exit()

        self.metals = []
        self.donors = {}

        self.sites = {}
        self.ligands_site = {}
  
        self.sitesSource = {}

   
    def findMetals(self):
        'Take out from pdb all the metal atoms found in list of metals'
        
        self.metals = [atom for atom in self.pdb.query('non-protein and model 1') if atom.elementType.upper().strip() in metalList]

        if not self.metals:
            raise Exception('Unable to find specified metal in the PDB file.')
            sys.exit()
        else:    
            for atom in self.metals:
                atom.beta = 50.00


    def findDonors(self): 
        
        not_donors = ['C','H','D']
        #distance = 3.0
        distance = 2.8
        atoms_around_metal = []
        
        ligands_names = {}
    
        for metal in self.metals:
            self.donors[metal] = []
            atoms_around_metal = self.pdb.query('model 1 & within {0} of'.format(distance), metal)
    
            for atom in atoms_around_metal:
                if ((atom.elementType.upper().strip() not in not_donors) and (atom.elementType.upper().strip() not in metalList)):
                    atom.beta = 40.00
                    self.donors[metal].append(atom)
                
                if not self.donors[metal]:
                    ligands_names[metal] = '-'
                else:
                    temp = []
                    for lig in self.donors[metal]:
                        temp.append("%s_%s_%s" % (lig.aa.strip(), lig.resid, lig.chain.strip()))
                    ligands_names[metal] = list(set(temp))
        
        return ligands_names


    def findSites(self, ligands_names):
        'Cluster metals into sites'
        
        maxDist = 3.5
        unassigned = [m for m in self.metals]

        site_numb = 0
        while unassigned:
            site_numb = site_numb + 1
            self.sites[site_numb] = []
            self.sites[site_numb].append(unassigned[0])
            unassigned.remove(unassigned[0])
            new_member = True
            while new_member:
                new_member = False
                for metal in unassigned:
                    for mt in self.sites[site_numb]:
                        distance = mt.distanceTo(metal)
                        ligands_in_common = list(set(ligands_names[metal]) & set(ligands_names[mt]))
                        if len(ligands_in_common) == 1 and ligands_in_common[0] == 'no_ligands':
                            ligands_in_common = []
                        if distance < maxDist or len(ligands_in_common) != 0:
                            self.sites[site_numb].append(metal)
                            unassigned.remove(metal)
                            new_member = True
                            break

  
    def findFirstSphere(self, userSites=None):
        
        if userSites:
            self.sites = userSites
            
        for site in self.sites:
            self.sitesSource[site] = []
            self.ligands_site[site] = []
            self.sitesSource[site].extend(self.sites[site]) #50.00 - metals
            
            for metal in self.sites[site]:
                donors_temp = [donor for donor in self.donors[metal]]  
            
            donors = list(set(donors_temp))

            atomsFromSameDonorResidue_temp = []
            for atom in donors:
                atoms_in_lig = self.pdb.query('model 1 and resid %s and chain %s' % (atom.resid, atom.chain.strip()))
                
                for atom in atoms_in_lig:
                    if ((atom in donors) or (atom in self.metals)):
                        continue
                    atom.beta = 30.00
                    atomsFromSameDonorResidue_temp.append(atom)

            atomsFromSameDonorResidues = list(set(atomsFromSameDonorResidue_temp))
            
            self.ligands_site[site].extend(donors) #40.00 - donors
            self.ligands_site[site].extend(atomsFromSameDonorResidues) #30.00 - ligands 1st sphere
            
            self.sitesSource[site].extend(donors)
            self.sitesSource[site].extend(atomsFromSameDonorResidues)


    def findSecondSphere(self, sites=None):
        
        if sites:
            self.sites = sites
        
        allAtomsOfSecondSphere_temp = []
        secondSphereDonors_temp = []
        
        for site in self.sites:    
            for atom in self.ligands_site[site]:
                temp = self.pdb.query('model 1 and within %s of' % (self.maxDist), atom)
    
                for atm in temp:
                    if ((atm in self.ligands_site[site]) or (atm in self.metals)):
                        continue
                    if (atm.elementType.upper().strip() != 'H' and atm.elementType.upper().strip() != 'D'):
                            atm.beta = 20.00
                            secondSphereDonors_temp.append(atm)

            secondSphereDonors = list(set(secondSphereDonors_temp))
            self.sitesSource[site].extend(secondSphereDonors)
            
            for atm in secondSphereDonors:
                atoms_in_sec = self.pdb.query('model 1 and resid %s and chain %s' % (atm.resid, atm.chain.strip()))
                
                for atom in atoms_in_sec:
                    if ((atom in secondSphereDonors) or (atom in self.metals)):
                        continue
                    atom.beta = 10.00
                    allAtomsOfSecondSphere_temp.append(atom)

            allAtomsOfSecondSphere = list(set(allAtomsOfSecondSphere_temp))
            self.sitesSource[site].extend(allAtomsOfSecondSphere)

  
    def findSitesInPDB(self):
        
        try:
            self.findMetals()
            ligands_names = self.findDonors()
            self.findSites(ligands_names)
        except:
            raise Exception('Unable to extract sites')


    def getSites(self):
        
        self.findFirstSphere()
        self.findSecondSphere(self.sites)
        
        metalSites = []
        
        for site in self.sites:
            try:
                path_to_site_file = self.dumpSiteToFile(site)
                inSite = Protein(path_to_site_file)
            except:
                raise Exception('Unable to open %s' % self.code)
            
            try:
                metalSite = MetalSite(inSite)
            except Exception:
                os.remove(path_to_site_file)
                continue
            else:
                metalSites.append(metalSite)			
                os.remove(path_to_site_file)
                
        return metalSites


    def getSiteWithMetal(self, metalID):
        
        metalSites = []
        siteWithMetal = {}
        metalIsDetected = False
        
        for site in self.sites:
            for metal in self.sites[site]:
                if metal.resid == metalID:
                    metalIsDetected = True
                    siteWithMetal[site] = [m for m in self.sites[site]]
                    break
                else:
                    break
            if metalIsDetected:
                break
        
        self.findFirstSphere(siteWithMetal)
        self.findSecondSphere(siteWithMetal)
        
        path_to_site_file = self.dumpSiteToFile(site)
        
        try:
            inSite = Protein(path_to_site_file)
        except:
            raise Exception('Unable to open %s' % code)
        else:
            metalSite = MetalSite(inSite)
            metalSites.append(metalSite)			
            os.remove(path_to_site_file)        
        
        return metalSites
    
    
    def dumpSiteToFile(self, site):    
        
        metal_sites_dir = os.path.dirname(os.path.abspath(__file__))
        output_name = '%s_%s.site.pdb' % (self.code, site)
        mSitePath = os.path.join(metal_sites_dir, output_name)

        outFile = open(mSitePath, 'w')
        
        outList = []
        atomSet = sorted(self.sitesSource[site], key=attrgetter('idx'))
        for atom in atomSet:
            outList.append("%s\n" % atom.output())
        outFile.writelines(outList)
        
        outFile.close()
        
        return mSitePath


#===============================================================================
"""METHODS"""
#===============================================================================


def getSiteFromSiteFile(pathSite):

    name = os.path.basename(pathSite)
    try:
        inSite = Protein(pathSite)
    except:
        raise Exception('Unable to open %s' % name)
    else:
        site = MetalSite(inSite)
        
    return site


def getSitesFromPdbFile(pathPdb, metalID=None, maxDist=5.0):
    
    sourcePDB = SourcePDB(pathPdb, maxDist)
    try:
        sourcePDB.findSitesInPDB()
        if metalID:
            metalSites = sourcePDB.getSiteWithMetal(metalID)
        else:
            metalSites = sourcePDB.getSites()
    except Exception:
        sys.exit()

    return metalSites