## \example kernel/dependency_graph.py
# When trying to understand what is going on in \imp, it can often be useful
# to view the dependency graph, that is, the graph showing how various
# entities relate to one another. In it, an arrow leads from an IMP::Container
# or IMP::Particle to an IMP::Restraint if the IMP::Restraint reads from that
# container or particle. Similarly, an arrow leads from an IMP::Container or
# IMP::Particle to an IMP::ScoreState if the score state reads from it, and
# an arrow leads from an IMP::ScoreState to an IMP::Container or IMP::Particle
# if the score state updates the particle.
#
# The resulting pruned dependency graph is:
# \dotgraph{ \dot
# digraph pruned_dependencies {
# 0[label="Nup85-12"];
# 1[label="Fragment [213-218)"];
# 2[label="Nup145C-4"];
# 3[label="Atom k distance restraint 2"];
# 4[label="Nup1202-8"];
# 5[label="Model 0"];
# 6[label="Fragment [239-244)"];
# 7[label="Nup1332-10"];
# 8[label="Nup1201-12"];
# 9[label="Connectivity particles"];
# 10[label="Nup84-14"];
# 11[label="Atom k distance restraint 3"];
# 12[label="MovedSingletonContainer2"];
# 13[label="Nup1333-16"];
# 14[label="Atom k distance restraint 5"];
# 15[label="Atom k distance restraint 4"];
# 16[label="normalize rigid bodies"];
# 17[label="Atom k distance restraint 0"];
# 18[label="ListSingletonContainer2"];
# 19[label="Chain A rigid body"];
# 20[label="ConnectivityRestraint 0"];
# 21[label="MovedSingletonContainer1"];
# 22[label="ListSingletonContainer1"];
# 23[label="HarmonicUpperBoundSphereDistancePairScore2 and ConnectingPairContainer"];
# 24[label="HarmonicUpperBoundSphereDistancePairScore1 and ConnectingPairContainer"];
# 25[label="Sec13"];
# 26[label="Hierarchy EV"];
# 27[label="Chain A rigid body rigid body positions"];
# 28[label="rigid bodies list"];
# 29[label="Atom k distance restraint 1"];
# 30[label="HarmonicUpperBoundSphereDistancePairScore0 and ConnectingPairContainer"];
# 31[label="ConnectingPairContainer"];
# 32[label="ConnectingPairContainer"];
# 33[label="Chain A rigid body"];
# 34[label="ListSingletonContainer0"];
# 35[label="ConnectivityRestraint 1"];
# 36[label="MovedSingletonContainer0"];
# 37[label="Chain A rigid body rigid body positions"];
# 38[label="ConnectingPairContainer"];
# 39[label="Seh1"];
# 40[label="Hierarchy EV particles"];
# 16->19 ;
# 16->33 ;
# 27->6 ;
# 27->39 ;
# 37->25 ;
# 37->1 ;
# 4->3 ;
# 2->3 ;
# 11->5 ;
# 17->5 ;
# 23->5 ;
# 14->5 ;
# 35->5 ;
# 15->5 ;
# 30->5 ;
# 24->5 ;
# 29->5 ;
# 26->5 ;
# 3->5 ;
# 20->5 ;
# 13->11 ;
# 10->11 ;
# 2->12 ;
# 18->12 ;
# 25->14 ;
# 2->14 ;
# 1->14 ;
# 33->14 ;
# 0->15 ;
# 6->15 ;
# 19->15 ;
# 39->15 ;
# 28->16 ;
# 8->17 ;
# 4->17 ;
# 13->20 ;
# 7->20 ;
# 22->21 ;
# 10->21 ;
# 2->23 ;
# 38->23 ;
# 31->24 ;
# 10->24 ;
# 0->26 ;
# 13->26 ;
# 8->26 ;
# 7->26 ;
# 4->26 ;
# 2->26 ;
# 6->26 ;
# 19->26 ;
# 10->26 ;
# 40->26 ;
# 1->26 ;
# 33->26 ;
# 19->27 ;
# 0->29 ;
# 2->29 ;
# 0->30 ;
# 32->30 ;
# 21->31 ;
# 22->31 ;
# 10->31 ;
# 0->32 ;
# 36->32 ;
# 34->32 ;
# 25->35 ;
# 2->35 ;
# 10->35 ;
# 1->35 ;
# 33->35 ;
# 0->36 ;
# 34->36 ;
# 33->37 ;
# 12->38 ;
# 2->38 ;
# 18->38 ;
# }
# \enddot
# }

import IMP.atom
import IMP.container
import sys

IMP.setup_from_argv(sys.argv, "Example of dependency graphs")


def create_representation():
    m = IMP.Model()
    all = IMP.atom.Hierarchy.setup_particle(IMP.Particle(m))
    all.set_name("the universe")
    rs = []

    def create_protein(name, ds):
        h = IMP.atom.create_protein(m, name, 10, ds)
        # leaves = IMP.atom.get_leaves(h)
        all.add_child(h)
        r = IMP.atom.create_connectivity_restraint(
            [IMP.atom.Selection(c) for c in h.get_children()], 1)
        if r:
            rs.append(r)

    def create_protein_from_pdbs(name, files):
        def create_from_pdb(file):
            with IMP.SetLogState(IMP.NONE):
                t = IMP.atom.read_pdb(
                    IMP.get_example_path("data/" + file), m,
                    IMP.atom.ATOMPDBSelector())
            # IMP.atom.show_molecular_hierarchy(t)
            c = IMP.atom.Chain(IMP.atom.get_by_type(t, IMP.atom.CHAIN_TYPE)[0])
            if c.get_number_of_children() == 0:
                IMP.atom.show_molecular_hierarchy(t)
            # there is no reason to use all atoms, just approximate the pdb
            # shape instead
            s = IMP.atom.create_simplified_along_backbone(c,
                                                          10.0 / 2.0)
            IMP.atom.destroy(t)
            # make the simplified structure rigid
            rb = IMP.atom.create_rigid_body(s)
            rb.set_coordinates_are_optimized(True)
            return s
        if len(files) > 1:
            p = IMP.Particle(m)
            h = IMP.atom.Hierarchy.setup_particle(p)
            h.set_name(name)
            for i, f in enumerate(files):
                c = create_from_pdb(f)
                h.add_child(c)
                c.set_name(name + " chain " + str(i))
            r = IMP.atom.create_connectivity_restraint(
                [IMP.atom.Selection(c) for c in h.get_children()], 1)
            if r:
                rs.append(r)
        else:
            h = create_from_pdb(files[0])
            h.set_name(name)
        all.add_child(h)
    create_protein("Nup85", 570)
    create_protein("Nup84", 460)
    create_protein("Nup145C", 442)
    create_protein("Nup120", [0, 500, 761])
    create_protein("Nup133", [0, 450, 778, 1160])
    create_protein_from_pdbs("Seh1", ["seh1.pdb"])
    create_protein_from_pdbs("Sec13", ["sec13.pdb"])
    return (m, rs, all)


def create_restraints(m, rs, all):
    def add_connectivity_restraint(s):
        r = IMP.atom.create_connectivity_restraint(s, 1)
        rs.append(r)

    def add_distance_restraint(s0, s1):
        r = IMP.atom.create_distance_restraint(s0, s1, 0, 1)
        rs.append(r)
    evr = IMP.atom.create_excluded_volume_restraint([all])
    rs.append(evr)
    s0 = IMP.atom.Selection(hierarchy=all, molecule="Nup145C",
                            residue_indexes=range(0, 423))
    s1 = IMP.atom.Selection(hierarchy=all, molecule="Nup84")
    s2 = IMP.atom.Selection(hierarchy=all, molecule="Sec13")
    add_connectivity_restraint([s0, s1, s2])
    add_distance_restraint(
        IMP.atom.Selection(hierarchy=all, molecule="Nup145C",
                           residue_indexes=range(0, 423)),
        IMP.atom.Selection(hierarchy=all, molecule="Nup85"))
    add_distance_restraint(
        IMP.atom.Selection(hierarchy=all, molecule="Nup145C",
                           residue_indexes=range(0, 423)),
        IMP.atom.Selection(hierarchy=all, molecule="Nup120",
                           residue_indexes=range(500, 762)))
    add_distance_restraint(IMP.atom.Selection(hierarchy=all, molecule="Nup84"),
                           IMP.atom.Selection(
                               hierarchy=all, molecule="Nup133",
                               residue_indexes=range(778, 1160)))
    add_distance_restraint(IMP.atom.Selection(hierarchy=all, molecule="Nup85"),
                           IMP.atom.Selection(hierarchy=all, molecule="Seh1"))
    add_distance_restraint(
        IMP.atom.Selection(hierarchy=all, molecule="Nup145C",
                           residue_indexes=range(0, 423)),
        IMP.atom.Selection(hierarchy=all, molecule="Sec13"))


# now do the actual work
(m, rs, all) = create_representation()
create_restraints(m, rs, all)

sf = IMP.core.RestraintsScoringFunction(rs)

# we can get the full dependency graph for the whole model with all
# the restraints but it is pretty complex
dg = IMP.get_dependency_graph(m)
IMP.show_graphviz(dg)

# better thing to do is to get the "pruned" graph
pdg = IMP.get_pruned_dependency_graph(m)

IMP.show_graphviz(pdg)
