## \example domino/merge_tree.py
# The example shows how to generate and inspect a merge tree for use in Domino.

import IMP
import IMP.domino
import IMP.core
import IMP.container
import IMP.algebra
import sys

IMP.setup_from_argv(sys.argv, "merge tree")

IMP.set_log_level(IMP.TERSE)
m = IMP.Model()
# don't print messages about evaluation
m.set_log_level(IMP.SILENT)

bb = IMP.algebra.BoundingBox3D((0, 0, 0), (10, 10, 10))
allc = []
for i in range(0, 7):
    p = IMP.Particle(m)
    d = IMP.core.XYZR.setup_particle(
        p, IMP.algebra.Sphere3D(IMP.algebra.get_random_vector_in(bb), 1))
    allc.append(d.get_coordinates())
pst = IMP.domino.ParticleStatesTable()
ss = IMP.domino.XYZStates(allc)
for p in m.get_particle_indexes():
    pst.set_particle_states(m.get_particle(p), ss)
# generate a set of restraints based on the close pairs in this randomly
# chosen configuration
cp = IMP.core.GridClosePairsFinder()
cp.set_distance(1)
cps = cp.get_close_pairs(m, m.get_particle_indexes())

if len(cps) > 0:
    # one cannot create a container from an empty list
    acp = IMP.container.ListPairContainer(m, cps)
else:
    acp = IMP.container.ListPairContainer(m)
ps = IMP.core.SoftSpherePairScore(1)
r = IMP.container.PairsRestraint(ps, acp)

# compute the interaction graph based on all the restraints
ig = IMP.domino.get_interaction_graph([r],
                                      pst)
# generate a junction tree from the interaction graph
jt = IMP.domino.get_junction_tree(ig)
print(dir(jt))
print(type(jt))
# create a merge tree from the junction tree, this can be passed to
# IMP.domin.DominoSampler
mt = IMP.domino.get_merge_tree(jt)
s = pst.get_subset()
print(s, type(s))
