#    LigAlign is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    LigAlign is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# LigAlign: Copyright 2008 Abraham Heifets

import tkSimpleDialog
import tkMessageBox
from pymol import cmd
from pymol import models
import pymol
import sys, random
import mwmatching

colors = [
"aquamarine",
#"black",
"blue",
"bluewhite",
"br0",
"br1",
"br2",
"br3",
"br4",
"br5",
"br6",
"br7",
"br8",
"br9",
"brightorange",
"brown",
"carbon",
"chartreuse",
"chocolate",
"cyan",
"darksalmon",
"dash",
"deepblue",
"deepolive",
"deeppurple",
"deepsalmon",
"deepsalmon",
"deepteal",
"density",
"dirtyviolet",
"firebrick",
"forest",
"gray",
"green",
"greencyan",
"grey",
"hotpink",
"hydrogen",
"lightblue",
"lightmagenta",
"lightorange",
"lightpink",
"lightteal",
"lime",
"limegreen",
"limon",
"magenta",
"marine",
"nitrogen",
"olive",
"orange",
"oxygen",
"palecyan",
"palegreen",
"paleyellow",
"pink",
"purple",
"purpleblue",
"raspberry",
"red",
"ruby",
"salmon",
"sand",
"skyblue",
"slate",
"smudge",
"splitpea",
"sulfur",
"teal",
"tv_blue",
"tv_green",
"tv_orange",
"tv_red",
"tv_yellow",
"violet",
"violetpurple",
"warmpink",
"wheat",
"white",
"yellow",
"yelloworange"
]

if False:
    cmd.reinitialize()
    cmd.load("testfiles/1UOF_pen2.pdb")
    cmd.create("a","1UOF_pen2")
    cmd.translate([10,0,0], "a")
    cmd.label("all", "ID")
    cmd.zoom("all")
    a=cmd.get_model("a")
    choice_atom = a.atom[0]
    cmd.select("temp", "all and id 2139")
    for n in cmd.get_names("all"):
        # test which names contain atoms w/ the appropriate id
        for atom in cmd.get_model(n).atom:
            print "n:",n, choice_atom.id, "vs", atom.id, "==", (choice_atom is atom), (choice_atom == atom), cmp(choice_atom, atom)
            if choice_atom == atom:
                print "found it!", choice_atom, atom
    #print a.atom[0].id, a.atom[0].
    cmd.iterate("temp", 'print "name:%s resn:%s resi:%s chain:%s alt:%s elem:%s q:%s b:%s segi:%s id:%s idx:%s" % (name, resn, resi, chain, alt, elem, q, b, segi, ID, index)')
    raise Exception

def init_defaults():
    if not pymol.stored.__dict__.has_key("ligalign_settings"):
        pymol.stored.ligalign_settings = {}    
#    if not cmd.setting.SettingIndex.hasattr("ligalign_"):
#        setattr(cmd.setting.SettingIndex, "ligalign_", True)
    pymol.stored.ligalign_settings["cutoff"] = 0.1
    pymol.stored.ligalign_settings["max_splits"] = None
    pymol.stored.ligalign_settings["min_fragment_size"] = 4
    pymol.stored.ligalign_settings["fingerprint_scoring"] = (lambda (a,b): arbitrary_scoring(a,b))
    pymol.stored.ligalign_settings["distant_residue_distance"] = 5.0
init_defaults()

def ligalign_set(*args):
    if not pymol.stored.__dict__.has_key("ligalign_settings"):
        init_defaults()

    if not args:
        for setting in pymol.stored.ligalign_settings.items():
            print "%s = %s", setting
    else:
        if len(args) == 2:
            # unwrap into setting and value, then set
            pymol.stored.ligalign_settings[args[0]] = float(args[1]) # pymol.stored.ligalign_settings[args[1]]
            pass
        else:
            print "I didn't quite understand that."
            raise QuietException
    pass
cmd.extend("ligalign_set", ligalign_set)

def random_color():
    return random.choice(colors)

def selection_from_atomlist(atomlist):
    return " or ".join(["(%s and id %d)" % (selection_name, id) for (selection_name, id) in atomlist])
def atomlist_from_selection(selection):
    return cmd.identify(selection,1)

def clean_sel(sel):
    cmd.alter(sel, "name='C1'")
    cmd.alter(sel, "resi='1'")
    cmd.alter(sel, "segi=''")
    cmd.alter(sel, "chain=''")
    cmd.alter(sel, "resn=''")
    cmd.alter(sel, "b=0.0")
    cmd.alter(sel, "q=0.0")

def rmsd_on_models(m1, m2, mapping, group):
    if len(group) == 0:
        return 0.0
    else:
        atoms1 = []
        atoms2 = []
        for index in group:
            # translate index to id
            id = -1
            for a in m1.atom:
                if a.index == index:
                    id = a.id
                    if id in mapping:
                        atoms1.append(a)
                        break
            for a in m2.atom:
                if id in mapping and a.id == mapping[id]:
                    atoms2.append(a)
                    break
        
        coords1 = [a.coord for a in atoms1]
        coords2 = [a.coord for a in atoms2]
        try: 
            (t1, t2, rot_mtx, rmsd) = models.fit(coords1, coords2)
        except Exception:  # TODO: figure out what's causing the errors, internal to the PyMOL code
            rmsd = 9999.0
        return rmsd

def align_by_ligand((baseline, junk_mapping), to_align): # selection, list of selections
    if not pymol.stored.__dict__.has_key("ligand_mappings"):
        pymol.stored.ligand_mappings = {}    
    maps = [(baseline, junk_mapping)]
    for (x,mapping) in to_align:
        if not mapping:
            mapping = build_mapping(x, baseline) #pymol.stored.ligand_mappings.setdefault((x,baseline), build_mapping(x, baseline))
        maps.append((x,mapping))
        # convert to pair list
        pairs = []
        x_model = cmd.get_model(x)
        for atom in x_model.atom:
            if atom.id in mapping:
                k,v = atom.id , mapping[atom.id]
                pairs.append("(%s and id %s)" % (x, k))                
                pairs.append("(%s and id %s)" % (baseline, v))
#        print "pair_fitting: ",pairs
        cmd.pair_fit(*pairs)
#        for i in range(0, len(pairs), 2):
#            name = "foo_%d" % i
#            cmd.distance(name, pairs[i], pairs[i+1])
#            cmd.show("dashes", name)
#            cmd.hide("labels", name)
#            print name, ":", pairs[i], pairs[i+1]
    for (name, mapping) in maps:
        pymol.stored.ligand_mappings[name] = (baseline, mapping)    # cache the mappings for possible future display
    return maps

def build_protein_specific_active_sites(distance=4):
    psas = ["(("+p+" within %d" % distance +" of "+ p +"_ligand) and not "+p+"_ligand)" for p in pymol.stored.protein_names]
    return psas

def build_ligand_union(ligand_selects):
    ligand_union = "ligand_union"
    cmd.select( "ligand_union", "((" + " or ".join(["(%s)" % l for l in ligand_selects]) + "))")
    cmd.deselect()
    return ligand_union

def build_active_site(ligand_selects, distance=4):
    # build ligand-union selection
    # then build AS selection
    ligand_union = "(" +" or ".join(["(%s)" % l for l in ligand_selects])+ ")"
    AS_select = "((all within %d"  % distance +" of "+ligand_union+") and not "+ligand_union+")"
    AS_select =  "(" +" or ".join(["(%s)" % l for l in build_protein_specific_active_sites(distance)])+ ")"
    cmd.select("as_select", AS_select)
    cmd.deselect()
    return AS_select

def draw_active_site(AS_select):
    cmd.show("spheres", AS_select)
    cmd.zoom(AS_select)

def viz(setting="1", ligands="ligand_union", active_site="as_select", full_residues=True, entire_protein=False, small_balls=False, distance=4):
    # global view effects
    # clear
    cmd.hide("everything", "all")
#    cmd.hide("lines", "all")
#    cmd.hide("everything", "all and resn HOH")
    if full_residues=="True" or full_residues=="true":
        tmp_selection_name = active_site
        full_active_site = "full_residues"
        cmd.select(full_active_site, "byres "+active_site)
        cmd.deselect()
        active_site = full_active_site
    if entire_protein=="True" or entire_protein=="true":
        for p in pymol.stored.protein_names:
            cmd.show("lines", p)
        cmd.zoom()
    if small_balls=="True" or small_balls=="true":
        # viz 2) ligand-sticks, AS-little balls
        cmd.show("spheres", active_site)
        cmd.show("lines", active_site)
        cmd.set("sphere_scale", 0.25, active_site)
        pass

    if setting == "1":
        # viz 1) ligand-sticks, AS-balls
        cmd.show("sticks", ligands)
        cmd.show("spheres", active_site)
        cmd.set("sphere_scale", 1.0, active_site)
        pass
    elif setting == "2":
        # viz 2) ligand-sticks, AS-little balls
        cmd.show("sticks", ligands)
        cmd.show("spheres", active_site)
        cmd.show("lines", active_site)
        cmd.set("sphere_scale", 0.25, active_site)
        pass
    elif setting == "3":
        # viz 3) ligand-sticks, as-closed surface transparent
        cmd.create("as_surface", active_site)
        cmd.show("sticks", ligands)
        cmd.show("spheres", "as_surface")
        cmd.show("lines", "as_surface")
        cmd.set("sphere_scale", 0.25, "as_surface")
        cmd.show("surface", "as_surface")
        cmd.set("transparency", 0.4)
        pass
    elif setting == "4":
        # viz 4) ligand-sticks, as-closed surface opaque
        cmd.create("as_surface", active_site)
        cmd.show("sticks", ligands)
        cmd.show("spheres", "as_surface")
        cmd.show("lines", "as_surface")
        cmd.set("sphere_scale", 0.25, "as_surface")
        cmd.show("surface", "as_surface")
        cmd.set("transparency", 0.0)
        pass
    elif setting == "5":
        # viz 5) ligand-sticks, as-open surface
        cmd.show("sticks", ligands)
        cmd.show("surface", active_site)
        cmd.set("transparency", 0.0)
        pass
    elif setting == "6":
        # viz 6) ligand-balls, as-sticks
        cmd.show("spheres", ligands)
        cmd.show("sticks", active_site)
        cmd.set("sphere_scale", 1.0, active_site)
        pass
    elif setting == "7":
        # viz 8)
        cmd.show("sticks", ligands)
        psas = build_protein_specific_active_sites()
        for i in range(len(pymol.stored.protein_names)):
            p = psas[i]
            if full_residues == "True":
                p = "(byres "+p+")"
            p_name = pymol.stored.protein_names[i]
            l_name = pymol.stored.ligand_names[i]
            prot_surface = p_name + "_surface"
            cmd.create(prot_surface, p)
            cmd.show("surface", prot_surface)
            cmd.hide("spheres", prot_surface)
            color = random_color()
            print "coloring",l_name,color
            cmd.color(color, prot_surface)
            cmd.color(color, l_name)
        cmd.set("transparency", 0.4)
        pass
    elif setting == "8":
        # viz 9)
        cmd.show("sticks", ligands)
        psas = build_protein_specific_active_sites()
        for i in range(len(pymol.stored.protein_names)):
            p = psas[i]
            if full_residues == "True":
                p = "(byres "+p+")"
            p_name = pymol.stored.protein_names[i]
            l_name = pymol.stored.ligand_names[i]
            prot_surface = p_name + "_surface"
            cmd.create(prot_surface, p)
            cmd.show("surface", prot_surface)
            cmd.show("sticks",  prot_surface)
            cmd.set("stick_radius", 0.10, prot_surface)
            cmd.hide("spheres", prot_surface)
            color = random_color()
            print "coloring",l_name,color
            cmd.color(color, prot_surface)
            cmd.color(color, l_name)
            cmd.color(color, p_name)
#            cmd.show("lines", p_name)
        cmd.set("transparency", 0.4)
        cmd.set("two_sided_lighting", 1)
        pass
    else:
        cmd.hide("everything", "all")
        pass
    cmd.center(ligands)
    cmd.zoom(active_site)
cmd.extend("ligalign_viz", viz)


def draw_mapping(baseline, to_align, mapping):
        if not mapping:
            mapping = build_mapping(to_align, baseline) #pymol.stored.ligand_mappings.setdefault((x,baseline), build_mapping(x, baseline))
        x_model = cmd.get_model(to_align)
        name = "mapping_%s_to_%s" % (to_align, baseline)
        for atom in x_model.atom:
            if atom.id in mapping:
                k,v = atom.id , mapping[atom.id]
                cmd.distance(name, "(%s and id %s)" % (to_align, k), "(%s and id %s)" % (baseline, v))
                cmd.show("dashes", name)
                cmd.hide("labels", name)
#                print name, ":", "(%s and id %s)" % (to_align, k), "(%s and id %s)" % (baseline, v)
def ligalign_show_mapping(sel):
    if not pymol.stored.__dict__.has_key("ligand_mappings"):
        return # no mappings means 
    if sel not in pymol.stored.ligand_mappings:
        ligand = ("%s_ligand" % sel)
        if ligand in pymol.stored.ligand_mappings:
            sel = ligand
        else:
            print "I don't have any mapping for",sel
            return
    (baseline,mapping) = pymol.stored.ligand_mappings[sel]
    draw_mapping(baseline, sel, mapping)
cmd.extend("ligalign_show_mapping", ligalign_show_mapping)

def color_by_closeness(baseline, active_site):
    if not pymol.stored.__dict__.has_key("ligalign_settings"):
        pymol.stored.ligalign_settings = {}
    if not pymol.stored.ligalign_settings.has_key("distant_residue_distance"):
        pymol.stored.ligalign_settings["distant_residue_distance"] = 5.0
    if not pymol.stored.__dict__.has_key("ligalign_colors") or not pymol.stored.ligalign_colors:
        pymol.stored.ligalign_colors = []
        for i in range(10):
            name = "ligalign_color%d"%i
            cmd.set_color(name, [(0.1*i), 1.0-(0.1*i), 0.0])
            pymol.stored.ligalign_colors.append(name)
    def find_color_bucket(distance):
        max_buckets = len(pymol.stored.ligalign_colors)
        color = None
        for i in range(0, max_buckets): #max_buckets-1, 0, -1):
            max_dist = float(pymol.stored.ligalign_settings["distant_residue_distance"])
#            print "dist:", distance, "i", i, "bucket:", i * (max_dist / max_buckets)
            if distance < i * (max_dist / max_buckets):
                color = pymol.stored.ligalign_colors[i]
                break
        if not color:
            color = pymol.stored.ligalign_colors[max_buckets-1]
        print "color:", color
        return color
    def compute_space(model):
        # for each atom in active_site, figure out its residue
        residue_ids = dict([(atom.id, atom.resi) for atom in model.atom])
        # select from all an atom.id to find the residue?  Worry about same residue number in different chains
        residue_atoms = dict([(residue, [atom for atom in model.atom if atom.resi == residue]) for residue in dict([(r,True) for r in residue_ids.values()]).keys()])
        # for each residue, compute centroid (BUG TODO FIXME: this isn't actually the centroid)
        avg = lambda (index, atom_list): sum([atom.coord[index] for atom in atom_list])/float(len(atom_list))
#        print "residue_atoms:",residue_atoms.items()
        centroids = dict([(residue, [avg((i, atom_list)) for i in [0,1,2]]) for (residue, atom_list) in residue_atoms.items()])
        return  residue_ids, residue_atoms, centroids 
    baseline_model = cmd.get_model(baseline)
    active_site_model = cmd.get_model(active_site)
    baseline_residue_ids, baseline_residue_atoms, baseline_centroids = compute_space(baseline_model)
    residue_ids, residue_atoms, centroids = compute_space(active_site_model)
    # for each centroid in active_site, figure out closest neighbor in baseline
    distances = dict([(residue, min([((centroid[0]-b_c[0])**2+(centroid[1]-b_c[1])**2+(centroid[2]-b_c[2])**2) ** 0.5 for (b_r, b_c) in baseline_centroids.items()])) for (residue,centroid) in centroids.items()])
    # foreach residue in active_site, color by min-centroid-distance  USEDTOBE: for each atom in centroid, color by min-centroid-distance
    for (residue, atoms) in residue_atoms.items():
        color = find_color_bucket(distances[residue])
        cmd.color(color, "(%s and resi %s)"%(active_site, residue))
#        print  "cmd.color(color, '(%s and resi %s)')" %(active_site, residue)
#    print "centroids:",centroids
#    print "dist:", distances

    if False:
        cmd.fragment(name="methane", object="%d"%i)#, name=("%d"%i), pos=[i,0,0], color="c")
        cmd.translate([i,0,0], "%d"%i)
        cmd.color("red", "%d"%i)
        cmd.color("c%d"%i, "%d"%i)
        raise Exception
cmd.extend("ligalign_color", color_by_closeness)



def get_similar_ligands(ligand, num_to_return=3):
    if not pymol.stored.__dict__.has_key("ligand_similarities") or not pymol.stored.ligand_similarities:
        print "loading similarity database..."
        pymol.stored.ligand_similarities = {}
        for line in file(os.getcwd()+"/similarityDB/sim_85_to_1", "r"):  # Figure out where to find this file, robustly.
            words = line.split()
            if len(words) != 3: continue
            similars = pymol.stored.ligand_similarities.setdefault(words[0], [])
            similars.append(words[1])
            similars = pymol.stored.ligand_similarities.setdefault(words[1], [])
            similars.append(words[0])
    similars = pymol.stored.ligand_similarities.get(ligand, [])
    return similars[0:num_to_return]

def fetch_pdbs(names):
    if not pymol.stored.__dict__.has_key("protein_names"):
        pymol.stored.protein_names = []    
    for p in names:
        print "fetching: ",p
        error = cmd.fetch(p, p, async=0)
        if not error:
            if p not in pymol.stored.protein_names:
                pymol.stored.protein_names.append(p)
    cmd.deselect()

def find_ligands():
    contenders = {}
    sorted = {}
    hack = {}
    for p in pymol.stored.protein_names:
        # pull out the hetatm residues in this protein
        pymol.stored.resi = {}
        cmd.select("temp", "(%s and (hetatm or het) and not resn HOH)" % p)
        cmd.iterate("temp", "stored.resi[(resi,chain)]=(resi,chain)")
        for (resi,chain) in pymol.stored.resi.keys():
            if chain=='':
                select =  "%s and resi %s" % (p, resi)
            else:
                select =  "%s and chain %s and resi %s" % (p,chain,resi)
            cmd.select("temp", select)
            nAtoms = cmd.count_atoms("temp")
            # now filter out the ones that are too tiny for us to do much analysis on (<=5 atoms)
            if nAtoms <= 5: continue
            print "analyzing", select
            m = cmd.get_model("temp")
            fingerprint = build_fingerprint(m)
            mol_fingerprint = [0] * len(fingerprint[0])
            for f in fingerprint:
                mol_fingerprint = [(x+y) for (x,y) in zip(mol_fingerprint, f)]
            mol_fingerprint = [x/nAtoms for x in mol_fingerprint]
            res_list = contenders.setdefault(p, [])
            res_list.append(((resi,chain), nAtoms, mol_fingerprint))
        if p not in contenders:
            raise Exception("No suitable ligand in %s!  Maybe all of the HETATMs are less than six heavy atoms?" % p)
    cmd.delete("temp")

    original_query = pymol.stored.protein_names[0] # assumes the order of finding proteins has been preserved!
    if not original_query in contenders:
        raise QuietException("No suitable ligand in %s!" % original_query)
    differences = {}
    # for each ligand in the original query,
    for (query_residue, query_nAtoms, query_fingerprint) in contenders[original_query]:
        # build a list of closest-sized ligands in the other proteins
        differences[query_residue] = []
        for p in contenders: #need to add query-molecule in somewhere.  pymol.stored.protein_names[1:]:# don't get points for matching against other ligands in query molecule.  #
            ligand_list = contenders[p]
            diff_list = [((100-score_fingerprints(fingerprint,query_fingerprint)),abs(nAtoms-query_nAtoms), -nAtoms, (p,resi,chain)) for ((resi,chain),nAtoms,fingerprint) in ligand_list]
            diff_list.sort()
            differences[query_residue].append(diff_list[0]) # smallest difference

    # choose ligands w/ min error, breaking ties to larger ligands 
    score = []
    for residue in differences:
        # compute total error
        sum_error = 0
        sum_size_diff = 0
        sum_size = 0
        residues = []
        for (diff, size_diff, size, lig) in differences[residue]:
            sum_error += diff
            sum_size_diff  += size_diff
            sum_size  -= size
            residues.append(lig)
        score.append((sum_error, sum_size_diff, sum_size, residues))
    score.sort()
    (min_error, min_size_diff, max_size, ligands) = score[0]

    ligand_names = []
    # NEED LIGAND ORDER TO MATCH PROTEIN ORDER!
    for pname in pymol.stored.protein_names:
        for (p, resi, chain) in ligands:
            if p==pname:
                ligand_name = p+"_ligand"
                ligand_names.append(ligand_name)
                if chain=='':
                    select =  "%s and resi %s" % (p, resi)
                else:
                    select =  "%s and chain %s and resi %s" % (p,chain,resi)
                cmd.select(ligand_name, select)
    cmd.deselect()
    pymol.stored.ligand_names = ligand_names
    return ligand_names

def random_ligand():
    get_similar_ligands(None) # init ligand_similarities db
    return random.choice(pymol.stored.ligand_similarities.keys())

def ligalign(fragment_based_alignment, ligand_tuple):
    reset_stored()
    ligands = list(ligand_tuple)
    if len(ligands) == 0:
        ligands = [random_ligand()]
    if len(ligands) == 1:
        ligands.extend(get_similar_ligands(ligands[0]))
    print "fetching PDB files..."
    fetch_pdbs(ligands)
    if len(pymol.stored.protein_names) < 1:
        raise Exception("No PDB files found for '%s'" % (", ".join(ligands)))
    print "extracting ligands..."
    ln = [(x,None) for x in find_ligands()]
    if fragment_based_alignment:
        print "fracturing ligands..."
        ln = fracture_ligands(ln[0], ln[1:len(ln)])
    print "aligning..."
    ln = align_by_ligand(ln[0], ln[1:len(ln)])
    print "displaying..."
    ligand_union = build_ligand_union([x for (x,y) in ln])
    active_site  = build_active_site([x for (x,y) in ln])
#    cmd.hide("everything", "all")
#    draw_active_site(active_site)
    viz(setting="7", full_residues="True") 
    if False: #True:
        baseline,junk_mapping = ln[0]
        for (x,m) in ln[1:len(ln)]:
            draw_mapping(baseline, x, m)
    cmd.deselect()
cmd.extend("ligalign",lambda *x,**y: ligalign(False, list(x)))
cmd.extend("ligalign_fragments",lambda *x,**y: ligalign(True, list(x)))
def reset_stored():
    # TODO: delete selections?
    cmd.reinitialize()
    pymol.stored.protein_names = []
    pymol.stored.ligand_names  = []
    pymol.stored.best_hinges_cache = {}
    pymol.stored.ligalign_colors = []
cmd.extend("ligalign_reset_stored",reset_stored)

def get_pdbs(ligand):
    fetch_pdbs([ligand])
    sims = get_similar_ligands(ligand)
    fetch_pdbs(sims)

class Bondset:
    def __init__(self, from_m, to_m, mapp, c={}, f={}, o=None, hscache = {}, ancache = {}, clump=None ):
        self.model    = from_m
        self.closed   = c.copy() # Pymol distributed with python < 2.3!  Use hashtables to implement sets
        self.front = f.copy()
        self.mapping  = mapp  # remember, mapping is in terms of ids
        self.target_model = to_m
        if o:
            self.open = o.copy()
        else:
            self.open = dict([(a.index,True) for a in from_m.atom if a.index not in c])
        self.hinge_score_cache = hscache
        self.atom_neighbors_cache = ancache
        if clump:
            self.clumps = clump
        else:
            self.clumps = self.calculate_clumps()
    def find_cycles(self):
        if len(self.model.atom) == 0:
            return {}
        arbitrary_atom = self.model.atom[0].index
        return self.find_cycles_helper(arbitrary_atom, None, [])
    def find_cycles_helper(self, atom, parent, stack):
        cycles = []
        if atom in stack: # we've found a cycle!
            cycle = stack[stack.index(atom):]
            cycles.append(cycle)
        else:
            # Recurse!
            stack.append(atom)
            children = self.atom_neighbors(atom)
            for child in children:
                if child != parent:
                    cycles.extend(self.find_cycles_helper(child, atom, stack))
            stack.pop()
        return cycles
    def calculate_clumps(self):
        # depth-first search to find cycles
        cycles = self.find_cycles()
        # union together cycles into cyclical systems
        while True:
            break_out  = False
            all_merged = True
            # cycles is a list of lists of atoms 
            for cycle_index in range(len(cycles)): # for each cycle
                for atom in cycles[cycle_index]:   # for each atom in each cycle, 
                    for cycle2 in range(cycle_index+1, len(cycles)):  # check that it's not in some other cycle
                        if atom in cycles[cycle2]:  # if it is, we need to merge
                            # need to merge cycles
                            all_merged = False
                            break_out = True
                            merged_cycle = cycles[cycle_index]
                            merged_cycle.extend([a for a in cycles[cycle2] if a not in merged_cycle])
                            del cycles[cycle2]
                            del cycles[cycle_index]
                            cycles.append(merged_cycle) # and insert the merged-cycle back into our list of cycles
                            break
                    if break_out: break
                if break_out: break
            if all_merged: break

        for cycle in cycles:
            cycle.sort()
        clumps = dict([(atom, cycle) for cycle in cycles for atom in cycle])
        return clumps
    def atom_neighbors(self, atom):
        if atom not in self.atom_neighbors_cache:
            atom_index = [a.index for a in self.model.atom].index(atom)
            neighbors = []
            # unforunately, bonds only record indexes, not ids
            for b in self.model.bond:
                if b.index[0] == atom_index:
                    neighbors.append(b.index[1])
                elif b.index[1] == atom_index:
                    neighbors.append(b.index[0])
#            return [self.model.atom[n].index for n in neighbors]
            self.atom_neighbors_cache[atom]=[self.model.atom[n].index for n in neighbors]
        return self.atom_neighbors_cache[atom]
    def expansion_set(self, a):
        return self.clumps.get(a,[a])
    def expand_frontier(self, a):
        clumps = self.expansion_set(a)  # if atom isn't part of a cycle, just add the atom.
        for atom in clumps:
            self.front.pop(atom, None) # using pop instead of del to suppress KeyError if atom not in open
            self.open.pop(atom, None) # using pop instead of del to suppress KeyError if atom not in open
            self.closed[atom] = True 
        # add it's neighbors (which haven't already been closed)
            for n in self.atom_neighbors(atom):
                if n not in self.closed:
                    self.front[n] = True 
        return len(clumps)
    def copy(self):
        return Bondset(self.model, self.target_model, self.mapping, self.closed, self.front, self.open, self.hinge_score_cache, self.atom_neighbors_cache, self.clumps)
    def frontier(self):
        return self.front.keys()
    def remaining_atoms(self):
        return self.open.keys()
    def __str__(self):
        return "closed="+str(self.closed.keys())+", frontier="+str(self.front.keys())
    def check_remainder_connected(self):
        test = self.copy()
        remainder = test.remaining_atoms()
        if len(remainder) == 0:
            return True
        # choose arbitrary atom
        arbitrary_atom = remainder[0]
        # clean out frontier, leaving an aribtrarily chosen atom
        test.front = {arbitrary_atom:True}
        # take transitive closure over remaining atoms (from arbitrary starting point)
        while len(test.front) > 0:
            (arbitrary_atom,junk) = test.front.popitem()
            test.expand_frontier(arbitrary_atom)
        # if every atom hasn't been closed, there was something unreachable from our starting point
        remainder = test.remaining_atoms() 
        return len(remainder) == 0
    def free(self):
        return self.remaining_atoms()
    def hinge_score(self, group):
        # caching
        group.sort()
        key = tuple(group)
        if key not in self.hinge_score_cache:
            self.hinge_score_cache[key] = self.hinge_score_helper(group)
        return self.hinge_score_cache[key]
    def hinge_score_helper(self, group):
        rmsd = rmsd_on_models(self.model, self.target_model, self.mapping, group)
        return rmsd ** 2.0
    def compute_model_key(self):
        atoms1 = [a.id for a in self.model.atom]
        atoms2 = [a.id for a in self.target_model.atom]
        atoms1.sort()
        atoms2.sort()
        mapping_list = self.mapping.items()
        mapping_list.sort()
        return (tuple(atoms1),tuple(atoms2),tuple(mapping_list))
    def seeds(self):
        # return each leaf atom and one atom from each ring system
        
        # grab first atoms out of each ring system
        # there might be some redundancies, so do the usual wash through a set
        ring_atoms = dict([(c[0], True) for c in self.clumps.values() if c[0] not in self.closed]).keys()
        
        # find leaf atoms (atoms which only have one neighbor that is still open (or are alone, i.e. last atom standing))
        leaf_atoms = [a for a in self.free() if len([x for x in self.atom_neighbors(a) if x not in self.closed]) < 2]

        seed_list = list(ring_atoms)
        seed_list.extend(leaf_atoms)
#        print "seeds:", seed_list, "size:",len(seed_list), " rings=",ring_atoms," leafs=",leaf_atoms
        return seed_list

def find_reasonable_splits(atomset):
    global debug_call_count
    cutoff = pymol.stored.ligalign_settings["cutoff"]
    max_splits = pymol.stored.ligalign_settings["max_splits"]
    splits = 0
    print "Attempting to use",splits,"hinges on",len(atomset.model.atom),"atoms..."
    score, size, groups = best_hinges(atomset,splits)
    score = score ** 0.5
    print "... tested",debug_call_count,"fragmentations.  Best gave an RMSD=",score,"with groups=",groups
    debug_call_count=0
    while True:
        splits += 1
        print "Attempting to use",splits,"hinges..."
        c, s, g = best_hinges(atomset,splits)
        c = c ** 0.5
        print "... tested",debug_call_count,"fragmentations.  Best gave an RMSD=",c,"with groups=",g,"size=",s
        assert len(atomset.model.atom) == s
        debug_call_count=0
        diff = (score - c)
        if diff < score * cutoff :
            print "Fragmenting into",splits,"pieces gave less than a", cutoff*100,"% improvement.  Let's use the fewer groups:  RMSD=",score,"with groups=",groups
            break
        elif max_splits and splits >= max_splits:
            print "Let's only use",splits,"fragments."
            break
        else:
            score, size, groups = c, s, g
    return score, groups

def fracture_ligands((baseline, junk_mapping), ligands):
    if not pymol.stored.__dict__.has_key("ligand_mappings"):
        pymol.stored.ligand_mappings = {}    # cache the mappings for possible future display
    prefractured_mappings = []
    new_ligands = [baseline]
    new_ligands.extend(ligands)
    return_vals = [(baseline, junk_mapping)]
    baseline_model = cmd.get_model(baseline)
    for (ligand, mapping) in ligands:
        ligand_model = cmd.get_model(ligand)
        print "Building mapping for",ligand,"and",baseline
        if not mapping:
            mapping = build_mapping(ligand, baseline)
        return_vals.append((ligand,mapping))

        bondset = Bondset(ligand_model, baseline_model, mapping)
        score, groups = find_reasonable_splits(bondset)
        nAtoms = cmd.count_atoms(ligand)
        nAtoms2 = sum([len(g) for g in groups],0)
        print "size of original ligand (",ligand,")=",nAtoms," and size of groups=",nAtoms2
        assert nAtoms == nAtoms2
        fragment_id = 0
        for g in groups:
            fragment_id += 1
            # create a duplicate ligand and protein for this group
            protein_name = ligand[:-7] # chops '_ligand' off ligand name to regenerate protein name
            new_protein_name = "%s_%d" % (protein_name, fragment_id) 
            frag_name = "%s_ligand" % new_protein_name
            cmd.create(new_protein_name, protein_name)
            frag = "(%s and (%s))" % (new_protein_name, " or ".join(["(index %d)" % (index) for index in g]))
            cmd.select(frag_name, frag)
            # insert them in the ligand/protein lists, # TODO: in the correct spots
            pymol.stored.protein_names.append(new_protein_name)
            pymol.stored.ligand_names.append(frag_name)
            new_ligands.append(frag_name)
            prefractured_mappings.append((frag_name,mapping))
            return_vals.append((frag_name,mapping))
#            pymol.stored.ligand_mappings[(frag_name,baseline)] = mapping
        # TODO: delete the original ligand 
    cmd.deselect()
    for (name, mapping) in return_vals:
        pymol.stored.ligand_mappings[name] = (baseline, mapping)    # cache the mappings for possible future display
    return return_vals #prefractured_mappings


debug_call_count = 0
debug_calls = 0
def best_hinges(atomset, splits, indent="", bound=sys.maxint):
    global debug_calls
    if not pymol.stored.__dict__.has_key("best_hinges_cache"):
        pymol.stored.best_hinges_cache = {}   
    # make hash-key
    free = atomset.free()
    free.sort()
    key = (splits, tuple(free), atomset.compute_model_key())
    # caching here
    if key not in pymol.stored.best_hinges_cache:
        answer = best_hinges_helper(atomset, splits, indent, bound)
        pymol.stored.best_hinges_cache[key] = answer
    return pymol.stored.best_hinges_cache[key] 
def best_hinges_helper(atomset, splits, indent, bound_score):
    global debug_call_count
    indent += "  "
    if bound_score < 0.0001:
        return None # someone's found an exact match, and we're just getting whacked by floating point precision
    elif splits == 0: 
        debug_call_count += 1 
        # no more splits, put all the remaining atoms into one group
        group = atomset.free()
        score = atomset.hinge_score(atomset.free())
        if score > bound_score:
            return None
        else:
            return score, len(atomset.free()), [atomset.free()]
    else:
        if len(atomset.free()) == 0: 
            group = atomset.free()
            debug_call_count += 1 
            # no more atoms, return an empty group
            return 0.0, 0, []
        else:
            # the option to use an extra hinge can only make the alignment better, not worse, so initialize the score to the score of one split less
            best = best_hinges(atomset, splits-1, indent, bound_score) # can't pass down bound b/c less splits on same atomset is harder!
            if not best: # can't beat bound
                best_score  = bound_score
                best_size   = None
                best_groups = None
            else:
                (best_score, best_size, best_groups) = best
                # a best could exist because it was calculated and cached at a looser bound,
                # which has since tightened, so we need to ensure that we can still use it.
                if bound_score <= best_score:
                    best_score  = bound_score
                    best_size   = None
                    best_groups = None
            for (group,rest) in reasonable_groups(atomset, indent):#, duplicates, work_reduction):
                group_score = atomset.hinge_score(group)
                group_size  = len(group)
                debug_call_count += 1 
                print indent, "#:",debug_call_count,"MSD:",best_score,"group:",group," group_score=",group_score," size=",group_size #," and the rest is", rest, " rest.free:",rest.free(), " len(rest.free):",len(rest.free()), " with ",splits," total splits"
                if ((group_score*group_size) / len(atomset.free())) >= best_score:
                    # we can tell the best_score will stay best, even if the error for 'rest' is 0
                    continue
                else:
                    candidate = best_hinges(rest, splits-1, indent, best_score)
                    if not candidate:
                        continue
                    (candidate_score,candidate_size,candidate_groups) = candidate
                    # combine
                    assert candidate_size == len(rest.free())
                    size  = (candidate_size + group_size)
#                    print indent, "candidate=",candidate
                    assert size == len(atomset.free())
                    score = ((candidate_score * candidate_size) + (group_score * group_size)) / size  #from unwrapping of RMSD
                    if score < best_score: # if scores are tied, use fewer hinges.  i.e, < not <= !
                        best_score  = score
                        best_size   = size
                        best_groups = list(candidate_groups)
                        best_groups.append(group)
            if best_score < bound_score:
                return (best_score, best_size, best_groups)
            else:
                return None
                    
def reasonable_groups(atomset, indent=""):
    min_fragment_size = int(pymol.stored.ligalign_settings["min_fragment_size"]) # small ligands tell us very little about the binding protein
    duplicates={}
    work_reduction={}
    for seed in atomset.seeds():
        seedgroup = list(atomset.expansion_set(seed))
        new_atomset = atomset.copy()
        new_atomset.front = {} # clear out old frontier
        new_atomset.expand_frontier(seed)
        if new_atomset.check_remainder_connected():  # remainder still connected
            nAtoms =  len(atomset.free())
            max_part_size = math.ceil(nAtoms/2.0) # 20 -> 10, 21 -> 11
#            for part_size in range(max_part_size-1, min_fragment_size-1, -1): #range(min_fragment_size,max_part_size):
            for part_size in range(min_fragment_size,max_part_size):
                for (group, rest) in extend_to_size(seedgroup, new_atomset, part_size, part_size, work_reduction, indent):
                    group.sort()
                    if tuple(group) not in duplicates:
                        duplicates[tuple(group)] = True
                        if rest.check_remainder_connected():
                            yield group,rest
                        else:
                            pass

# TODO: we can terminate fragment expansion early, if the number of disconnected components we've generated exceeds the size we have left to spend (if we have more size to spend, we could maybe still end up connecting them all).  Need to track the number of disconnected components but, since we take all rings in one piece, the number of disconnected components is equal to the size of the frontier!
def extend_to_size(seedgroup, atomset, size, max_size, work_reduction, indent=""):
    # make cache key: (size, frontier, group, atomset)
    front_dict = dict([(tuple(atomset.expansion_set(atom)),True) for atom in atomset.frontier()])
    front = front_dict.keys()
    front.sort()
    free = atomset.free()
    free.sort()
    group2 = list(seedgroup)
    group2.sort()
    #key = (size, max_size, tuple(front), tuple(free), tuple(group2), atomset.compute_model_key())
    key = (size, tuple(front), tuple(free), tuple(group2), atomset.compute_model_key())

    # caching here
    if key not in work_reduction:
        answer = extend_to_size_helper(seedgroup, atomset, size, max_size, work_reduction, indent)
        work_reduction[key] = answer
    return work_reduction[key]

def extend_to_size_helper(seedgroup, atomset, size, max_size, work_reduction, indent=""):
#    indent += "  "
    if size<=0:
        return [(seedgroup, atomset)]
    else:
        # remove duplicates
        unduplicated = {}
        for (group, rest) in extend_to_size(seedgroup, atomset, size-1, max_size, work_reduction, indent):
            group.sort()
            unduplicated[tuple(group)] = rest

        expanded_groups = []
        for (group_t, rest) in unduplicated.items():
            group = list(group_t)
            if len(group) > size:
                expanded_groups.append((group, rest))
            else:
                front_dict = dict([(tuple(rest.expansion_set(atom)),atom) for atom in rest.frontier()]) # remove duplicates
                for exp_set,atom in front_dict.items():
                    new_atomset = rest.copy()
                    new_atomset.expand_frontier(exp_set[0])
                    if True or (len(new_atomset.frontier()) <= max_size):
#                    if (len(new_atomset.frontier()) <= max_size):
                        expanded_groups.append((group+list(exp_set), new_atomset))
        return expanded_groups

# from Magnus Lie Hetland 
def floyd_warshall(weights, n):
    d = {-1: weights}
    for k in range(n):
        d[k] = {}
        for i in range(n):
            for j in range(n):
                if  d[k-1][i,k]==sys.maxint or d[k-1][k,j]==sys.maxint: 
                    second_sum = sys.maxint
                else: 
                    second_sum = d[k-1][i,k] + d[k-1][k,j] 
                d[k][i,j] = min(d[k-1][i,j], second_sum)            
    return d[n-1]

def weights_from_bonds(bonds,n):
    weights = {}
    # default values
    for i in range(n):
        weights[i,i] = 0
        for j in range(n):
            weights.setdefault((i,j), sys.maxint)
    # build symmetric connectivities
    for b in bonds:
        assert len(b.index) == 2
        i, j = b.index[0], b.index[1]
        weights[i,j] = 1
        weights[j,i] = 1
    return weights

def connectivity(model):
    nAtoms =  len(model.atom)
    weights = weights_from_bonds(model.bond, nAtoms)
    connections = floyd_warshall(weights, nAtoms)
    return connections

# assumes the model is connected
def build_fingerprint(model):
    nAtoms = len(model.atom)
    connections = connectivity(model)
    fingerprint = []
    counts = {}
    for x in range(nAtoms):
        counts[x] = {}
        for y in range(nAtoms):
            distance = connections[x,y]
            counts[x][distance] = counts[x].get(distance, 0)+1
        fingerprint.append([])
        for y in range(nAtoms): # can't be further apart than if the atoms were in a linear chain
            distance = counts[x].get(y, 0)
            if distance == 0:
                break
            fingerprint[x].append(distance)
    return fingerprint

def score_fingerprints(a,b):
    return pymol.stored.ligalign_settings["fingerprint_scoring"]((a,b))
    
def arbitrary_scoring(a,b):
    # A variety of scoring functions are possible.  Here are a few.
    if True:
        # weighted diff
        edit_dist   = [abs(x-y) for (x,y) in zip(a,b)] # truncates to shorter list.  Is this correct or should we pad with 0?  Probably pad, but this is easier for now
        multipliers = [4,4,4,3,3,3,2,2,2,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1] # for sensitivity to further away atoms, extend this list.  Also, these weights are completely arbitrary.
        weighted    = [x*y for (x,y) in zip(edit_dist, multipliers)]
        return 10000-sum(weighted, 0)
    elif False:
        # exponentially weighted diff
        edit_dist = [abs(x-y) for (x,y) in zip(a,b)]#[0:4] # truncates to local area.  Is this correct or should we pad with 0?  Probably pad, but this is easier for now
        val = 0.0
        mult = 100000000000.0 # we'll get precision errors w/ very large ligands, the kind where differences in atoms connectivity get below floating point precision
        for x in edit_dist:
            val += mult * x
            mult *= 0.1
        return 1000000000000.0-val
    elif False:
        # proportional to number of edits, but doesn't care how far away they are.
        edit_dist = [abs(x-y) for (x,y) in zip(a,b)][0:4] # truncates to local area.
        return 100 - sum(edit_dist,0)
    else:
        # hamming
        edit_dist = [1 for (x,y) in zip(a,b) if x!=y][0:4] # truncates to local area. 
        return 100 - sum(edit_dist,0)

def get_atom_id(model, atom_num):
    return model.atom[atom_num].id

def build_mapping(sel1, sel2):
    if not pymol.stored.__dict__.has_key("molecule_mappings"):
        pymol.stored.molecule_mappings = {}    
    if (sel1, sel2) in pymol.stored.molecule_mappings:
        return pymol.stored.molecule_mappings[(sel1,sel2)]


    def score_maps(weights_set, nAtoms1, model1, model2, group):
        # convert set to tuples
        p_weights = [(a,b,weights_set[(a,b)]) for (a,b) in weights_set]
        
        # mwmatching requires all verticies to be in one contiguous labelling,
        # so append the atoms numbers together, call it, then pull it back apart
        # matchings = mwmatching.maxWeightMatching(weights)
        matchings = mwmatching.maxWeightMatching(p_weights)
        one_onto_two = [matchings[i]-nAtoms1 for i in range(nAtoms1)]
        #    two_onto_one = [matchings[i] for i in range(nAtoms1,nAtoms1+nAtoms2)]

        # convert to index from ids
        mapping = {}
        for i in range(len(one_onto_two)):
            if one_onto_two[i] < 0: continue
            mapping[model1.atom[i].id] = model2.atom[one_onto_two[i]].id

        score = rmsd_on_models(model1, model2, mapping, group)
        return mapping,score

    model1  = cmd.get_model(sel1)
    model2  = cmd.get_model(sel2)
    finger1 = build_fingerprint(model1)
    finger2 = build_fingerprint(model2)
    nAtoms1 = len(model1.atom)
    nAtoms2 = len(model2.atom)
    group = [a.index for a in model1.atom]

    weights = []
    foo_weights = {}
    for i in range(nAtoms1):
        for j in range(nAtoms1, nAtoms1+nAtoms2):
            index2 = j-nAtoms1
            weight = score_fingerprints(finger1[i], finger2[index2])
            # include a bonus if the types of the atoms match.  Don't want this to be infinitely weighted, b/c we may want to align variants of ligands together
            if model1.atom[i].symbol == model2.atom[index2].symbol:
                weight += 4  # arbitrary, but chosen to be same as weight for distance 1
            i_weights = foo_weights.setdefault(i, {})
            i_nodes    = i_weights.setdefault(weight*10, [])
            i_nodes.append(j)
            weights.append((i,j,weight))

    if False:
        for i in foo_weights:
            print "\n\ni=",i
            foo_w_k = foo_weights[i].keys()
            foo_w_k.sort()
            foo_w_k.reverse()
            for w in foo_w_k:
                print "w=", w, foo_weights[i][w]
        raise Exception

    # for each node i, for each node j in a set of nodes that all have equal weight,
    # build a weighting that 
    p_weights_list = [[]]
    if False:
        for i_prime in foo_weights:
            for i in foo_weights:
                foo_w_k = foo_weights[i].keys()
            foo_w_k.sort()
            foo_w_k.reverse()        
            most_important_weights = foo_w_k[0:1] # take the first (and second?) heaviest weighted 
        #        print "i=",i," weights=", [(x, foo_weights[i][x]) for x in foo_w_k]
            for weight in foo_weights[i]:
                weight_class = foo_weights[i][weight]
                wc = []
                if weight in most_important_weights and i==i_prime:
                    for j in weight_class:
                        w = []
                # make j the strongest in this weight class
                        for node in weight_class:
                            if node == j:
                                tup = (i, node, weight+1)
                            else:
                                tup = (i, node, weight)
                            w.append(tup)
                        wc.append(w)
                else:
                    wc.append([(i,node,weight) for node in weight_class])
#            print "   weight=",weight,"wc=",wc
                new_list = []
                for p in p_weights_list:
                    for w in wc:
                        l = list(p)
                        l.extend(w)
                        new_list.append(l)
            p_weights_list = new_list #[list(p).extend(w) for p in p_weights_list for w in wc]


    p_weights_list = [[]]
    weights_set = {}
    weights_to_check = {}
    for i in foo_weights:
        # first, initialize the weights
        for weight in foo_weights[i]:
            for node in foo_weights[i][weight]:
                weights_set[(i,node)] = weight

        # second, segregate the list into the ones where there's only a single maximally-weighted neighbor, and those where there's several
        foo_w_k = foo_weights[i].keys()
        foo_w_k.sort()
        foo_w_k.reverse()        
        weight_class = foo_weights[i][foo_w_k[0]]
        if len(weight_class) != 1:
            weights_to_check[(i, foo_w_k[0])] = weight_class


    best_mapping = None
    decided_settings = []
    for (i, weight), nodes in weights_to_check.items():
        best_node, best_score = None, None
        for node in nodes:
            # make node the prefered partner (break ties in favor of node)
            weights_set[(i,node)] = weight+1
            (mapping,score) = score_maps(weights_set, nAtoms1, model1, model2, group)

            if not best_score or score < best_score:
                best_score   = score
                best_node    = node
                best_mapping = mapping

            # reset it back
            weights_set[(i,node)] = weight           
#            print "Mapping with score=",best_score,":", best_mapping
        # now fix the prefered partner to be the best found
        decided_settings.append((i,best_node,weight+1))

    for (a,b,c) in decided_settings:
        weights_set[(a,b)] = c
    (mapping,score) = score_maps(weights_set, nAtoms1, model1, model2, group)

    # cache mappings
#    print "Mapping with score=",best_score,":", best_mapping
    pymol.stored.molecule_mappings[(sel1, sel2)] = best_mapping #mapping
    return best_mapping #mapping # list of pairs of atom selections

def debug(sel):
    if True:
        cmd.iterate(sel, 'print "name:%s resn:%s resi:%s chain:%s alt:%s elem:%s q:%s b:%s segi:%s id:%s idx:%s" % (name, resn, resi, chain, alt, elem, q, b, segi, ID, index)')
cmd.extend("ligalign_debug_selection", debug)

#cmd.do("ligalign 1xbb, 1opj")
#cmd.do("ligalign_fragments 1xbb, 1opj")
#cmd.do("ligalign_color 1xbb_surface, 1opj_surface")
if False: # set true to profile code
    import hotshot, hotshot.stats
    prof = hotshot.Profile("ligalign.prof")
#profile.run('cmd.do("ligalign_fragments 1gzf, 2qwx")')
#profile.run('ligalign(True, ("1gzf","2qwx"))')
    prof.runcall(ligalign, True, ("1gzf", "2qwx"))
    prof.close()
    stats = hotshot.stats.load("ligalign.prof")
    stats.strip_dirs()
    stats.sort_stats('time', 'calls')
    stats.print_stats(50)

print "LigAlign loaded."
