##===============================================================================
#     This file is part of TEMPy.
#     
#     TEMPy is a software designed to help the user in the manipulation 
#     and analyses of macromolecular assemblies using 3D electron microscopy maps. 
#     
#     Copyright 2010-2014 TEMPy Inventors and Birkbeck College University of London.
#                          The TEMPy Inventors are: Maya Topf, Daven Vasishtan, 
#                           Arun Prasad Pandurangan, Irene Farabella, Agnel-Praveen Joseph,
#                          Harpal Sahota
# 
# 
#     TEMPy is available under Public Licence.
#     
#     Please cite your use of TEMPy in published work:
#     
#     Vasishtan D, Topf M. (2011) J Struct Biol 174:333-343. Scoring functions for cryoEM density fitting.
#
#===============================================================================

from TEMPy.ProtRep_Biopy import *
from TEMPy.StructureBlurrer import StructureBlurrer
#from EMMap import *
#from MapParser import *
#from PDBParser import *
#from StructureBlurrer import *
#from Vector import *
#from VQ import *
#from Quaternion import *


class Assembly:

    def __init__(self, structList):
        self.structList = structList
        self.initMapList = []
        self.mapList = []
        
    def build_maps(self, resolution, template_map, sig_coeff=0.356):
        sb = StructureBlurrer()
        for x in self.structList:
            self.mapList.append(sb.gaussian_blur(x, resolution, template_map, sig_coeff))
            self.initMapList.append(self.mapList[-1].copy())

    def randomise_structs(self, max_trans, max_rot, v_grain=30, rad=False):
        for x in self.structList:
            x.randomise_position(max_trans, max_rot, v_grain, rad)

    def randomise_structs_and_maps(self, max_trans, max_rot, v_grain=30, rad=False):
        if len(self.mapList) != len(self.structList):
            print 'Maps not built yet'
        else:
            for x in range(len(self.structList)):
                com = self.structList[x].CoM.copy()
                rx,ry,rz,ra,tx,ty,tz = self.structList[x].randomise_position(max_trans, max_rot, v_grain, rad, verbose=True)
                self.mapList[x] = self.mapList[x].rotate_by_axis_angle(rx, ry, rz, ra, com)
                self.mapList[x] = self.mapList[x].translate(tx,ty,tz)

    def reset_structs(self):
        for x in self.structList:
            x.reset_position()

    def reset_maps(self):
        for x in range(len(self.mapList)):
            self.mapList[x] = self.initMapList[x].copy()

    def reset_all(self):
        self.reset_maps()
        self.reset_structs()

    def move_map_and_prot_by_aa(self, index, rx, ry, rz, ra, tx, ty, tz):
        com = self.structList[index].CoM.copy()
        self.structList[index].rotate_by_axis_angle(rx, ry, rz, ra, tx, ty, tz)
        self.mapList[index] = self.mapList[index].rotate_by_axis_angle(rx, ry, rz, ra, com)
        self.mapList[index] = self.mapList[index].translate(tx,ty,tz)

    def move_map_and_prot_by_euler(self, index, rx, ry, rz, tx, ty, tz):
        com = self.structList[index].CoM.copy()
        self.structList[index].rotate_by_euler(rx, ry, rz, 0, 0, 0)
	self.structList[index].translate(tx,ty,tz)
        self.mapList[index] = self.mapList[index].rotate_by_euler(rx, ry, rz, com)
        self.mapList[index] = self.mapList[index].translate(tx,ty,tz)

    def move_map_and_prot_by_mat(self, index, mat, tx, ty, tz):
        com = self.structList[index].CoM.copy()
	self.structList[index].rotate_by_mat(mat)
	self.structList[index].translate(tx,ty,tz)
	self.mapList[index] = self.mapList[index].rotate_by_matrix(mat, com)
        self.mapList[index] = self.mapList[index].translate(tx,ty,tz)

    def move_map_and_prot_by_quat(self, index, tx, ty, tz, q_param, mat):
        com = self.structList[index].CoM.copy()
        self.structList[index].rotate_by_quaternion(q_param)
	self.structList[index].translate(tx,ty,tz)
        self.mapList[index] = self.mapList[index].rotate_by_matrix(mat, com)
        self.mapList[index] = self.mapList[index].translate(tx,ty,tz)

    def combine_structs(self):
        if len(self.structList)>1:
            return self.structList[0].combine_structures(self.structList[1:])
        elif len(self.structList)==1:
            return self.structList[0]
        else:
            print 'No structures found'

    def combine_maps(self):
        if len(self.mapList)>1:
            newMap = self.mapList[0].copy()
            for x in self.mapList[1:]:
                newMap.fullMap += x.fullMap
            return newMap
        elif len(self.structList)==1:
            return self.mapList[0].copy()
        else:
            print 'No maps found'
    '''
    def combine_selected_maps(self, list_indx):
	if len(list_indx) > 1:
		newMap = self.mapList[list_indx[0]].copy()
		for i in list_indx[1:]:
			newMap.fullMap += self.mapList[i].fullMap
		return newMap
	elif len(list_indx) == 1:
		return self.mapList[list_indx[0]].copy
	else:
		print 'Invalid map list index'
    '''
	
    def make_VQ_points(self, threshold, noOfPoints, lap_fil, epochs=300):
        vq = []
        if len(self.mapList) > 0:
            for s in range(len(self.mapList)):
                vq.append(get_VQ_points(self.mapList[s], threshold, noOfPoints[s], epochs, None, lap_fil))
        return vq

    def make_sub_VQ_points(self, i, threshold, noOfPoints, lap_fil, epochs=300):
        vq = []
        if len(self.mapList) > 0:
        	vq.append(get_VQ_points(self.mapList[i], threshold, noOfPoints, epochs, None, lap_fil))
        return vq

    def make_subVQ_points(self, threshold, noOfPoints, lap_fil, epochs=300):
        vq = []
        vq.append(get_VQ_points(self.mapList[0], threshold, noOfPoints, epochs, None, lap_fil))
        return vq


    def write_all_to_files(self, templateName):
        for x in range(len(self.structList)):
            self.structList[x].write_to_PDB(templateName+str(x)+'.pdb')
            if len(self.mapList) > 0:
                self.mapList[x].write_to_MRC_file(templateName+str(x)+'.mrc')