#!/usr/bin/env python
# -*- coding: utf-8 -*-

from p3d.protein import Protein
import copy
import sys
import os
from operator import attrgetter

from atom import atomIsFromLigand, atomIsNeighbour, atomIsFromBackbone, atomIsFromSideChain, atomIsFromDNAorRNA, atomIsFromProtein
from atom import getOneLetterResudueCode
from constants import *

class MetalSite(object):
	"""
	Single metal site object shifted to the origin of coordinates 

    usage MetalSite(metals, donors, sphere)

	metals:   list of metal atoms
	donors:   list of donor atoms (bounded to a metal directly)
	sphere:   list of atoms populating approximate sphere around metals
	"""
	
	def __init__(self, rawSite=None):

		if rawSite != None:
			
			self.nameExt = rawSite.filename.split('/')[-1]
			self.name = rawSite.filename.split('/')[-1].split('.')[0]
				
			metals = [atom for atom in rawSite.atoms if (atom.beta == metalsTempFactor)]
			donors = [atom for atom in rawSite.atoms if (atom.beta == donorsTempFactor)]
			sphere = [atom for atom in rawSite.atoms if ((atom.beta != metalsTempFactor) and (atom.beta != donorsTempFactor))]
				
			if not metals:
				print '\nNo metals are found or file in a bad format.'
				raise Exception
			if not donors:
				print '\n %s has no ligands or file in a bad format. Site will not be processed.' % (self.name)
				raise Exception
			else:
				self.gcX, self.gcY, self.gcZ = self.getGeometricalCentre(metals)
				self.gcX_old, self.gcY_old, self.gcZ_old = self.gcX, self.gcY, self.gcZ
				
				self.metals = self.shiftCoordinates(metals)
				self.donorAtoms = self.shiftCoordinates(donors)
				self.approxSphereAtoms = self.shiftCoordinates(sphere)
				
				#dictionaries
				self.getAtomsForMatching()
				
				#lists
				self.getBackboneAndSidechain()
		
		else:
			self.name = 'none'
			self.metals = []
			self.donorAtoms = []
			self.approxSphereAtoms = []
		
	def getOriginalGeomCentre(self):
		return self.gcX_old, self.gcY_old, self.gcZ_old
	
	def geometricalCentre(self):
		'Calculates geometric centre of metals and returns coordinates of the geometric centre'
	
		gcX = 0.0
		gcY = 0.0
		gcZ = 0.0
			
		nAtoms = len(self.metals)
	
		gcX = float(sum([metal.x for metal in self.metals]))/nAtoms
		gcY = float(sum([metal.y for metal in self.metals]))/nAtoms
		gcZ = float(sum([metal.z for metal in self.metals]))/nAtoms
	
		return gcX, gcY, gcZ
	
	def getGeometricalCentre(self, metals):
		'Calculates geometric centre of metals and returns coordinates of the geometric centre'
	
		gcX = 0.0
		gcY = 0.0
		gcZ = 0.0
			
		nAtoms = len(metals)
	
		gcX = float(sum([metal.x for metal in metals]))/nAtoms
		gcY = float(sum([metal.y for metal in metals]))/nAtoms
		gcZ = float(sum([metal.z for metal in metals]))/nAtoms
	
		return gcX, gcY, gcZ
	
	def shiftCoordinates(self, atoms):
		'Shifts all atoms by coordinates of geometric center'

		for atom in atoms:          
			atom.x = atom.x - self.gcX
			atom.y = atom.y - self.gcY
			atom.z = atom.z - self.gcZ
		
		return atoms

	def getAtomsForMatching(self):
	
		self.backboneDonors = {}
		self.backboneNeighbours = {}
		
		self.sideChainDonors = {}
		self.sideChainNeighbours = {}
		
		for atom in self.approxSphereAtoms:
			
			if atomIsFromBackbone(atom):
				if atomIsFromLigand(atom):
					if(atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom)):
						handle = '{0}{1}{2}'.format(atom.chain.upper().strip(), '_', atom.resid)
						if handle not in self.backboneDonors.keys():
							self.backboneDonors[handle] = atom
				elif atomIsNeighbour(atom):
					if(atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom)):
						handle = '{0}{1}{2}'.format(atom.chain.upper().strip(), '_', atom.resid)
						if handle not in self.backboneNeighbours.keys():
							self.backboneNeighbours[handle] = atom
			
			elif atomIsFromSideChain(atom):
				if atomIsFromLigand(atom):
					if(atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom)):
						handle = '{0}{1}{2}'.format(atom.chain.upper().strip(), '_', atom.resid)
						if (handle in self.backboneDonors.keys() and handle not in self.sideChainDonors.keys()):
							self.sideChainDonors[handle] = atom
				elif atomIsNeighbour(atom):
					if(atomIsFromProtein(atom) or atomIsFromDNAorRNA(atom)):
						handle = '{0}{1}{2}'.format(atom.chain.upper().strip(), '_', atom.resid)
						if (handle in self.backboneNeighbours.keys() and handle not in self.sideChainNeighbours.keys()):
							self.sideChainNeighbours[handle] = atom

	def getBackboneAndSidechain(self):
		
		self.backbone = self.backboneDonors.items() + self.backboneNeighbours.items()
		self.sideChain = self.sideChainDonors.items() + self.sideChainNeighbours.items()
	
	def getBackbone(self, chain_id=None):
		
		backbone = []
		backbone.extend(self.backboneDonors.values())
		backbone.extend(self.backboneNeighbours.values())
		
		if chain_id:
			backbone = [a for a in backbone if a.chain.strip() == chain_id]
		
		return sorted(backbone, key=attrgetter('idx'))
	
	def getSequence(self, chain_id=None):
		
		backbone = self.getBackbone(chain_id)
		self.ligandsFragmentsMapping = []
		
		i=0
		sequence = []
		last_residue = False
		fragment_has_ligs = False
		fragment_ligands = []
		
		while i < len(backbone):
				
			fragment = getOneLetterResudueCode(backbone[i])
			curr_resid = backbone[i].resid
			
			if atomIsFromLigand(backbone[i]) and backbone[i].type.strip()=='ATOM':
				fragment_ligands.append(backbone[i])
			i += 1
			
			while (backbone[i].resid - curr_resid == 1):
				
				curr_resid = backbone[i].resid
				fragment += getOneLetterResudueCode(backbone[i])

				if atomIsFromLigand(backbone[i]) and backbone[i].type.strip()=='ATOM':
					fragment_ligands.append(backbone[i])
					fragment_has_ligs = True
				
				if i == len(backbone)-1:
					last_residue = True
					break
				i += 1
			
			sequence.append(fragment)
					
			if not fragment_has_ligs:
				self.ligandsFragmentsMapping.append(False)
			else:
				self.ligandsFragmentsMapping.append(fragment_ligands)
				fragment_ligands = []
				fragment_has_ligs = False
			
			if last_residue:
				break
			
		return sequence
		
	def getLigandsFragmentsMapping(self):
			
		return self.ligandsFragmentsMapping
	
	def getLigandsProximity(self, ligands=None, window=2):
		
		if not ligands:
			ligands = self.getEndogenousCaLigands()
		backbone = self.getBackbone()
		ligands_proximity = []
		
		for lig in ligands:
			lig_prox = []
			inds = [x-window for x in range(window)] + [x+1 for x in range(window)]
			for i in inds:
				neighbour_residue = next((res for res in backbone if ((lig.chain.strip() == res.chain.strip()) and ((lig.resid+i) == res.resid))), None)
				lig_prox.append(neighbour_residue)
			ligands_proximity.append(lig_prox)
				
		return ligands_proximity
		
	def getEndogenousDonors(self, chain_id=None):
		
		donors = []
		for atom in self.donorAtoms:
			if atomIsFromProtein(atom):
				if chain_id:
					if atom.chain.strip() != chain_id:
						continue
				donors.append(atom)
		donors = sorted(donors, key=attrgetter('resid'))
		return donors
	
	def getEndogenousCaLigands(self, chain_id=None):
		
		backbone = self.getBackbone(chain_id)
		donors_list = list(set(map(lambda d: d.resid, filter(lambda x: atomIsFromProtein(x), self.donorAtoms))))
		
		atoms = []
		for atom in backbone:
			if atomIsFromProtein(atom) and atom.atype.strip()=='CA' and atom.resid in donors_list:
				atoms.append(atom)
				
		return sorted(atoms, key=attrgetter('resid'))
	
	def getSecondSphere(self):

		ligs = []
		for atom in self.donorAtoms:
			if atomIsFromProtein(atom):
				lig_key='%s_%i_%s' % (atom.aa.strip(), atom.resid, atom.chain.strip())
				ligs.append(lig_key)
				
		second_sphere = []
		for atom in self.approxSphereAtoms:
			if atomIsFromProtein(atom) and atomIsFromBackbone(atom):
				key='%s_%i_%s' % (atom.aa.strip(), atom.resid, atom.chain.strip())
				if key not in ligs:
					second_sphere.append(atom)

		return second_sphere
	
	def writeToFile(self, filename):
		'Writes an ASCII text representation of an MetalSite object to a file' 
		
		# Opens output file for writing    
		outFile = open(filename, 'w')
		
		# Aggregates atoms in a single list and sorts atoms in the list by index number  
		atomSet = sorted(self.metals + self.donorAtoms + self.approxSphereAtoms, key=attrgetter('idx'))
		
		outList = []
		for atom in atomSet:
			if atom.elementType is "X":
				atom.elementType  = " "
			outList.append("%s\n" % atom.output())
				  
		outFile.writelines(outList)    
		outFile.close()

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

def multipleCoordination(site1, site2):
	if (len(site1.donorAtoms) > 1 and len(site2.donorAtoms) > 1):
		return True
	else:
		return False

def singleCoordination(site1, site2):
	if (len(site1.donorAtoms) == 1 and len(site2.donorAtoms) == 1):
		return True
	else:
		return False

def getCoordinatesOfAtomsInList(atmList):
	'Extracts coordinates of atoms'

	coordinates  = [[atom.x, atom.y, atom.z] for atom in atmList] 

	return coordinates

def equipSite(infile):
    
    try:                
        inSite = Protein(infile)
    except:
        raise Exception
    else:
        site = MetalSite(inSite)
   
    return site

def copySite(site, tempSite):
    """
    Copies a given site to a tempSite
    
    Returns copied site
    """
    name = site.name
        
    metals = []
    for atom in site.metals:                    
        tempAtom = copy.copy(atom)
        tempAtom.x = atom.x
        tempAtom.y = atom.y
        tempAtom.z = atom.z  
        metals.append(tempAtom)

    donorAtoms = []
    for atom in site.donorAtoms:                    
        tempAtom = copy.copy(atom)
        tempAtom.x = atom.x
        tempAtom.y = atom.y
        tempAtom.z = atom.z    
        donorAtoms.append(tempAtom)            

    approxSphereAtoms = []
    for atom in site.approxSphereAtoms:                    
        tempAtom = copy.copy(atom)
        tempAtom.x = atom.x
        tempAtom.y = atom.y
        tempAtom.z = atom.z    
        approxSphereAtoms.append(tempAtom)
   
    tempSite.name = name
    
    tempSite.metals = metals
    tempSite.donorAtoms = donorAtoms
    tempSite.approxSphereAtoms = approxSphereAtoms    
    
    tempSite.getAtomsForMatching()
    tempSite.getBackboneAndSidechain() 
  
    return tempSite
