#
# Container for molecular conformations
#

from __future__ import division
import Source.group, Source.determinants, Source.determinant, Source.ligand, Source.output, Source.coupled_groups, functools

class Conformation_container:
    def __init__(self, name='', parameters=None, molecular_container=None):
        self.molecular_container = molecular_container
        self.name=name
        self.parameters=parameters
        self.atoms = []
        self.groups = []
        self.chains = []
        self.current_iter_item = 0
        
        self.marvin_pkas_calculated = False
        self.non_covalently_coupled_groups = False
       
        return


    #
    # Group related methods
    #
    def extract_groups(self):
        """ Generates at list of molecular groups needed for calculating pKa values """
        for atom in self.get_non_hydrogen_atoms():
            # has this atom been checked for groups?
            if atom.groups_extracted == 0:
                self.validate_group(Source.group.is_group(self.parameters, atom))
            # if the atom has been checked in a another conformation, check if it has a 
            # group that should be used in this conformation as well 
            elif atom.group:
                self.validate_group(atom.group)

        return 

    def additional_setup_when_reading_input_file(self):
        
        # if a group is coupled and we are reading a .propka_input file, 
        # some more configuration might be needed

        # print coupling map
        map = Source.output.make_interaction_map('Covalent coupling map for %s'%self,
                                                 self.get_covalently_coupled_groups(),
                                                 lambda g1,g2: g1 in g2.covalently_coupled_groups)
        print(map)

        # check if we should set a common charge centre as well
        if self.parameters.common_charge_centre:
            self.set_common_charge_centres()
                
        return
    
    def set_common_charge_centres(self):
        for system in self.get_coupled_systems(self.get_covalently_coupled_groups(), Source.group.Group.get_covalently_coupled_groups):
            # make a list of the charge centre coordinates
            all_coordinates = list(map(lambda g: [g.x, g.y, g.z], system))
            # find the common charge center
            ccc = functools.reduce(lambda g1,g2: [g1[0]+g2[0], g1[1]+g2[1], g1[2]+g2[2]], all_coordinates)
            ccc = list(map(lambda c: c/len(system), ccc))
            # set the ccc for all coupled groups in this system
            for g in system:
                [g.x, g.y, g.z] = ccc
                g.common_charge_centre = True
        return




    def find_covalently_coupled_groups(self):
        """ Finds covalently coupled groups and sets common charge centres if needed """
        for group in [ group for group in self.groups if group.titratable]:

            # Find covalently bonded groups
            bonded_groups = self.find_bonded_titratable_groups(group.atom, 1, group.atom)

            # couple groups
            for cg in bonded_groups:
                if cg in group.covalently_coupled_groups:
                    continue
                if cg.atom.sybyl_type == group.atom.sybyl_type:
                    group.couple_covalently(cg)
                    
        # check if we should set a common charge centre as well
        if self.parameters.common_charge_centre:
            self.set_common_charge_centres()

        # print coupling map
        map = Source.output.make_interaction_map('Covalent coupling map for %s'%self,
                                                 #self.get_titratable_groups(),
                                                 self.get_covalently_coupled_groups(),
                                                 lambda g1,g2: g1 in g2.covalently_coupled_groups)
        print(map)

            
        return


    def find_non_covalently_coupled_groups(self, verbose=False):
        # check if non-covalent coupling has already been set up in an input file
        if len(list(filter(lambda g: len(g.non_covalently_coupled_groups)>0, self.get_titratable_groups())))>0:
            self.non_covalently_coupled_groups = True
            
        Source.coupled_groups.nccg.identify_non_covalently_coupled_groups(self,verbose=verbose)

        # re-do the check
        if len(list(filter(lambda g: len(g.non_covalently_coupled_groups)>0, self.get_titratable_groups())))>0:
            self.non_covalently_coupled_groups = True

        return


    def find_bonded_titratable_groups(self, atom, no_bonds, original_atom):
        res = set()
        for ba in atom.bonded_atoms:
            # skip the original atom
            if ba == original_atom:
                continue
            # check if this atom has a titratable group
            if ba.group and ba.group.titratable and no_bonds <= self.parameters.coupling_max_number_of_bonds:
                res.add(ba.group)
            # check for titratable groups bonded to this atom
            if no_bonds < self.parameters.coupling_max_number_of_bonds:
                res |= self.find_bonded_titratable_groups(ba,no_bonds+1, original_atom)

        return res


    def validate_group(self, group):
        """ Checks if we want to include this group in the calculations """
        # Is it recognized as a group at all?
        if not group:
            return

        # Other checks (include ligands, which chains etc.)

        # if all ok, accept the group
        self.accept_group(group)
        return


    def accept_group(self, group):
        # set up the group
        group.parameters=self.parameters
        group.setup()

        # and store it
        self.groups.append(group)
        return


    #
    # pka calculation methods
    #

    def calculate_pka(self, version, options):
        print('\nCalculating pKas for',self)
                
        # calculate desolvation
        for group in self.get_titratable_groups()+self.get_ions():
            version.calculate_desolvation(group)

        # calculate backbone interactions
        Source.determinants.setBackBoneDeterminants(self.get_titratable_groups(), self.get_backbone_groups(), version)

        # setting ion determinants
        Source.determinants.setIonDeterminants(self, version)

        # calculating the back-bone reorganization/desolvation term
        version.calculateBackBoneReorganization(self)

        # setting remaining non-iterative and iterative side-chain & Coulomb interaction determinants
        Source.determinants.setDeterminants(self.get_sidechain_groups(), version=version, options=options)

        # calculating the total pKa values
        for group in self.groups: group.calculate_total_pka()
           
        # take coupling effects into account
        penalised_labels = self.coupling_effects()

        if self.parameters.remove_penalised_group and len(penalised_labels)>0:
            print('Removing penalised groups!!!')

            for g in self.get_titratable_groups():
                g.remove_determinants(penalised_labels)

            # re-calculating the total pKa values
            for group in self.groups: group.calculate_total_pka()


        return


    def coupling_effects(self):
        #
        # Bases: The group with the highest pKa (the most stable one in the 
        #        charged form) will be the first one to adopt a proton as pH 
        #        is lowered and this group is allowed to titrate.
        #        The remaining groups are penalised
        #
        # Acids: The group with the highest pKa (the least stable one in the
        #        charged form) will be the last group to loose the proton as 
        #        pH is raised and will be penalised.
        #        The remaining groups are allowed to titrate.
        #
        penalised_labels = []
        
        for all_groups in self.get_coupled_systems(self.get_covalently_coupled_groups(), 
                                                   Source.group.Group.get_covalently_coupled_groups):

            # check if we should share determinants
            if self.parameters.shared_determinants:
                self.share_determinants(all_groups)

            # find the group that has the highest pKa value
            first_group = max(all_groups, key=lambda g:g.pka_value)

            # In case of acids
            if first_group.charge < 0:
                first_group.coupled_titrating_group = min(all_groups, key=lambda g:g.pka_value)
                penalised_labels.append(first_group.label) # group with the highest pKa is penalised

            # In case of bases
            else:
                for a in all_groups:
                    if a == first_group:
                        continue # group with the highest pKa is allowed to titrate...
                    a.coupled_titrating_group = first_group 
                    penalised_labels.append(a.label) #... and the rest is penalised

        return penalised_labels


    def share_determinants(self, groups):
        
        # make a list of the determinants to share
        types = ['sidechain','backbone','coulomb']
        for type in types:
            # find maximum value for each determinant
            max_dets = {}
            for g in groups:
                for d in g.determinants[type]:
                    # update max dets
                    if d.group not in max_dets.keys():
                        max_dets[d.group] = d.value
                    else:
                        max_dets[d.group] = max(d.value, max_dets[d.group], key= lambda v: abs(v))

             # overwrite/add maximum value for each determinant
            for det_group in max_dets.keys():
                new_determinant = Source.determinant.Determinant(det_group, max_dets[det_group])
                for g in groups:
                    g.set_determinant(new_determinant,type)


        return


    def get_coupled_systems(self, groups, get_coupled_groups):
        """ This generator will yield one covalently coupled system at the time """
        groups = set(groups)        
        while len(groups)>0:
            # extract a system of coupled groups ...
            system = set()
            self.get_a_coupled_system_of_groups(groups.pop(), system, get_coupled_groups)
            # ... and remove them from the list
            groups -= system

            yield system

        return

    
    def get_a_coupled_system_of_groups(self, new_group, coupled_groups, get_coupled_groups):
        coupled_groups.add(new_group)
        [self.get_a_coupled_system_of_groups(c, coupled_groups, get_coupled_groups) for c in get_coupled_groups(new_group) if c not in coupled_groups]
        return


    #
    # Energy/summary-related methods
    #
    def calculate_folding_energy(self, pH=None, reference=None):
        ddg = 0.0
        for group in self.groups:
            #print('Folding energy for %s at pH %f: %f'%(group,pH,group.calculate_folding_energy(self.parameters, pH=pH, reference=reference)))
            ddg += group.calculate_folding_energy(self.parameters, pH=pH, reference=reference)

        return ddg

    def calculate_charge(self, parmaeters, pH=None):
        unfolded = folded = 0.0
        for group in self.get_titratable_groups():
            unfolded += group.calculate_charge(parmaeters, pH=pH, state='unfolded')
            folded   += group.calculate_charge(parmaeters, pH=pH, state='folded')

        return unfolded,folded


    #
    # conformation/bookkeeping/atom methods
    #

    def get_backbone_groups(self):
        """ returns all backbone groups needed for the pKa calculations """
        return [group for group in self.groups if 'BB' in group.type]

    def get_sidechain_groups(self):
        """ returns all sidechain groups needed for the pKa calculations """
        return [group for group in self.groups if ('BB' not in group.type\
                    and not group.atom.cysteine_bridge)]
    
    def get_covalently_coupled_groups(self):
        return [g for g in self.groups if len(g.covalently_coupled_groups)>0]

    def get_non_covalently_coupled_groups(self):
        return [g for g in self.groups if len(g.non_covalently_coupled_groups)>0]

    def get_backbone_NH_groups(self):
        """ returns all NH backbone groups needed for the pKa calculations """
        return [group for group in self.groups if group.type == 'BBN']

    def get_backbone_CO_groups(self):
        """ returns all CO backbone groups needed for the pKa calculations """
        return [group for group in self.groups if group.type == 'BBC']

    def get_groups_in_residue(self, residue):
        return [group for group in self.groups if group.residue_type == residue]

    def get_titratable_groups(self):
        return [group for group in self.groups if group.titratable]
    
    def get_titratable_groups_and_cysteine_bridges(self):
        return [group for group in self.groups if group.titratable or group.residue_type == 'CYS']

    def get_acids(self):
        return [group for group in self.groups if (group.residue_type in self.parameters.acid_list
                                                   and not group.atom.cysteine_bridge)]

    def get_backbone_reorganisation_groups(self):
        return [group for group in self.groups if (group.residue_type in self.parameters.backbone_reorganisation_list
                                                   and not group.atom.cysteine_bridge)]

    def get_ions(self):
        return [group for group in self.groups if group.residue_type in self.parameters.ions.keys()]

    def get_group_names(self, list):
        return [group for group in self.groups if group.type in list]


    def get_ligand_atoms(self):
        return [atom for atom in self.atoms if atom.type=='hetatm']

    def get_heavy_ligand_atoms(self):
        return [atom for atom in self.atoms if atom.type=='hetatm' and atom.element != 'H']

    def get_chain(self,chain):
        return [atom for atom in self.atoms if atom.chainID != chain]


    def add_atom(self, atom):
        #print(self,'adding',atom)
        self.atoms.append(atom)
        if not atom.conformation_container:
            atom.conformation_container = self
        if not atom.molecular_container:
            atom.molecular_container = self.molecular_container
        
        # store chain id for bookkeeping
        if not atom.chainID in self.chains:
            self.chains.append(atom.chainID)
            
        return 

    def copy_atom(self, atom):
        new_atom  = atom.makeCopy()
        self.atoms.append(new_atom)
        new_atom.conformation_container = self

        return 

    def get_non_hydrogen_atoms(self):
        return [atom for atom in self.atoms if atom.element!='H']


    def top_up(self, other):
        """ Tops up self with all atoms found in other but not in self """
        for atom in other.atoms:
            if not self.have_atom(atom):
                self.copy_atom(atom)
        return
    
    def have_atom(self, atom):
        res = False
        for a in self.atoms:
            if a.residue_label == atom.residue_label:
                res = True
                break
        return res

    def find_group(self, group):
        for g in self.groups:
            if g.atom.residue_label == group.atom.residue_label:
                if g.type == group.type:
                    return g
        return False

    
    def set_ligand_atom_names(self):
        for atom in self.get_ligand_atoms():
            Source.ligand.assign_sybyl_type(atom)
        return
            


    def __str__(self):
        return'Conformation container %s with %d atoms and %d groups'%(self.name,len(self),len(self.groups))

    def __len__(self):
        return len(self.atoms)

    def sort_atoms(self):
        # sort the atoms ...
        self.atoms.sort(key=self.sort_atoms_key)
        # ... and re-number them
        for i in range(len(self.atoms)):
            self.atoms[i].numb = i+1

        return

    def sort_atoms_key(self, atom):
        key = ord(atom.chainID)*1e7 
        key += atom.resNumb*1000
        if len(atom.name) > len(atom.element):
            key += ord(atom.name[len(atom.element)])
            #print(atom,ord(atom.name[len(atom.element)]), '|%s||%s|'%(atom.name,atom.element))
        return key