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

import os
import sys
sys.path.append(os.path.join(os.path.expanduser('~'), 'development'))

import glob, shutil

from operator import attrgetter
from nomenclature import metalList

from metalSite import MetalSite
from transformation import Rotations
from transformation import applyRotationMatrix2
from sequence import getSequenceOfAlignment_web, writeSequenceToFile

from helpers.utilites import askToOverwriteIt

class ScoreReport(object):
    'Collection of data for a score reports'

    def __init__(self, site1, site2):

        self.getNames(site1, site2)
        
        self.rmsd=None
        
        self.p1 = None
        self.p2 = None
        self.p3 = None
        
        self.totalScore = None
        self.percentage = None
        
        self.metals1=''
        self.metals2=''
        self.ligands1=''
        self.ligands2=''
        
        for atom in site1.metals:    
            self.metals1 += '%s_%s_%s, ' % (atom.aa, atom.resid, atom.chain)
        self.metals1 = self.metals1.strip().strip(',')
        
        for atom in site2.metals:    
            self.metals2 += '%s_%s_%s, ' % (atom.aa, atom.resid, atom.chain)
        self.metals2 = self.metals2.strip().strip(',')
        
        for atom in site1.donorAtoms:    
            self.ligands1 += '%s_%s_%s, ' % (atom.aa, atom.resid, atom.chain)
        self.ligands1 = self.ligands1.strip().strip(',')

        for atom in site2.donorAtoms:    
            self.ligands2 += '%s_%s_%s, ' % (atom.aa, atom.resid, atom.chain)
        self.ligands2 = self.ligands2.strip().strip(',')
        
        self.initial_score = None
        self.min_initial_score = None
        self.max_initial_score = None
        self.threshold_percentage = None
        
    def getNames(self, site1, site2):
        
        self.name1 = site1.name
        self.name2 = site2.name 
    
    def setScores(self, rmsd, p1, p2, p3, totalScore, percentage):
        
        self.rmsd = rmsd

        self.p1 = p1
        self.p2 = p2
        self.p3 = p3
        
        self.totalScore = totalScore
        self.percentage = percentage

def storeAlignment(site1, site2, outputPath):

    if not os.path.exists(outputPath):
        try:
            os.mkdir(outputPath)
        except OSError:
            print 'Alignment is not complete. Check if you have permissions to create directories.'
            sys.exit()

    rootDir = os.path.join(outputPath, '%s_vs_%s' % (site1.name, site2.name))
    try:
        os.mkdir(rootDir)
    except:
        print '\nOutput directory "%s" already exists.\nDo you want to overwrite it? (y/n)\n\nWarning! Deleting a directory will also delete all contents within the directory.\n' % (sitesDirRoot)
        overwriteDir = askToOverwriteIt()

        if overwriteDir: #The directory is to be overwritten
            shutil.rmtree(rootDir)
            os.mkdir(rootDir)
        if not overwriteDir:
            sys.exit()

    sitesDir = os.path.join(rootDir, 'sites')
    os.mkdir(sitesDir)

    writeSiteToFile(sitesDir,'%s_query.site.pdb' % (site1.name), site1)
    writeSiteToFile(sitesDir, '%s_target.site.pdb' % (site2.name), site2)
    visualizeAlignment(sitesDir)

    return rootDir

def storeRotations(root_path, site1, sit2):

    name1 = site1.name.split('.')[0]
    name2 = site2.name.split('.')[0]
    
    folder = name1+'_vs_'+name2
    path = os.path.join(root_path, folder)
    if not os.path.exists(path):
        os.mkdir(path)

    writeSiteToFile(path, name1, site1)
    writeSiteToFile(path, name2, site2)
    
    tempSite = MetalSite()
    
    rotations = Rotations(site1, site2)
    poses = rotations.getPoses()
    
    i=0
    for pose in poses:
        i = i+1
        rotM = pose.rotationMatrix
        name = name2 + '_rot_' + str(i)
        site = applyRotationMatrix2(site2, tempSite, rotM)
        writeSiteToFile(path, name, site)

def reportAlignment(site1, site2, scoreReport, outputPath, maxDist):
	
	scorePath = os.path.join(outputPath, 'score.txt')
	writeScoresToFile(scoreReport, scorePath)
	
	sequencePath = os.path.join(outputPath, 'sequence.txt')
	writeSequenceToFile(site1, site2, sequencePath, maxDist)

def writeScoresToFile(scoreReport, path):
    
    out_file = open(path, 'w')
    
    name1 = '%s_%s' % (scoreReport.name1, scoreReport.metals1)
    fst_line = '{0:<10}{1:<25}\n'.format('Name 1:', name1)
    out_file.write(fst_line)
    name2 = '%s_%s' % (scoreReport.name2, scoreReport.metals2)
    snd_line = '{0:<10}{1:<25}\n'.format('Name 2:', name2)
    out_file.write(snd_line)
    
    out_file.write('\n')
    
    rmsd_line = '{0:<27}{1:<25}\n'.format('RMSD over aligned region:', scoreReport.rmsd)
    out_file.write(rmsd_line)
    
    out_file.write('\n')

    p1_line = '{0:<27}{1:<25}\n'.format('Relative coverage term:', scoreReport.p1)
    out_file.write(p1_line)
    p2_line = '{0:<27}{1:<25}\n'.format('Chemical similarity term:', scoreReport.p2)
    out_file.write(p2_line) 
    p3_line = '{0:<27}{1:<25}\n'.format('Continuity term:', scoreReport.p3)
    out_file.write(p3_line)
    
    out_file.write('\n')
    
    total_line = '{0:<27}{1:<25}\n'.format('Total score:', scoreReport.totalScore)
    out_file.write(total_line)
    
    out_file.close()

def visualizeAlignment(path):
    
    fileName = 'visualisation.pml'    
    files = os.listdir(path)       
    filePath = os.path.join(path, fileName)
            
    out = open(filePath, 'w')
    
    outList = []
    for fname in glob.iglob( os.path.join(path, '*.pdb')):
        nm = 'load %s\n' % (fname.split('/')[-1])
        outList.append(nm)
  
    outList.append('deselect\n')
    outList.append('hide everything\nshow cartoon\n')
    outList.append('sele me,resn %s' % (metalList[0]))
    for i in xrange(1, len(metalList)):
        outList.append(' or resn %s' % (metalList[i]))           
    
    outList.append('\nshow spheres,me\nset sphere_scale, 0.65\ncolor red, me\ndeselect\n')
    out.writelines(outList)

    out.close()

def writeSiteToFile(path, name, site):
    'Writes an ASCII text representation of an MetalSite object to a file' 
    
    # Defines output name    
    trgPath = os.path.join(path, name)

    # Opens output file for writing    
    outFile = open(trgPath, 'w')
    
    # Aggregates atoms in a single list and sorts atoms in the list by index number  
    atomSet = sorted(site.metals + site.donorAtoms + site.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()
    
    return trgPath

def writePdbToFile(pdb, filename):
	
	atomSet = sorted(pdb.atoms, key=attrgetter('idx'))

	olist = []
	for atom in atomSet:
		if atom.elementType is "X":
			atom.elementType  = " "
		olist.append("%s\n" % atom.output())
	
	with open(filename, "w") as fl:
		fl.writelines(olist)

def reportSequenceToFile(site1, site2, outputPath, maxDist):

    sequencePath = os.path.join(outputPath, 'sequence')
    getSequenceOfAlignment_web(site1, site2, sequencePath, maxDist)
    
    sequencePath_txt = os.path.join(outputPath, 'sequence.txt')
    writeSequenceToFile(site1, site2, sequencePath_txt, maxDist)
        