from __future__ import print_function
import IMP
import IMP.test
import IMP.core
import IMP.container


class SingletonTestModifier(IMP.SingletonModifier):

    def __init__(self, k):
        IMP.SingletonModifier.__init__(self)
        self.k = k

    def do_show(self, fh):
        fh.write("Test Particle")

    def apply_index(self, m, a0, da=None):
        if m.get_has_attribute(self.k, a0):
            pass
        else:
            m.add_attribute(self.k, a0, 1)

    def get_version_info(self):
        return 1

    def do_get_inputs(self, m, pis):
        return [m.get_particle(i) for i in pis]

    def do_get_outputs(self, m, pis):
        return [m.get_particle(i) for i in pis]


class PairTestModifier(IMP.PairModifier):

    def __init__(self, k):
        IMP.PairModifier.__init__(self)
        self.k = k
        self.sm = SingletonTestModifier(k)

    def do_show(self, fh):
        fh.write("Test Particle")

    def apply_index(self, m, a0, da=None):
        self.sm.apply_index(m, a0[0])
        self.sm.apply_index(m, a0[1])

    def get_version_info(self):
        return 1

    def do_get_inputs(self, m, pis):
        return [m.get_particle(i) for i in pis]

    def do_get_outputs(self, m, pis):
        return [m.get_particle(i) for i in pis]


def particle_has_attribute(p, k):
    return p.has_attribute(k)


def particle_pair_has_attribute(p, k):
    return p[0].has_attribute(k) and p[1].has_attribute(k)


def particle_triplet_has_attribute(p, k):
    return (
        p[0].has_attribute(k) and p[1].has_attribute(
            k) and p[2].has_attribute(k)
    )

# This file is generated by the make-container script


class Tests(IMP.test.TestCase):

    """Tests for ClassnameContainer related objects"""

    def create_particle(self, m):
        p = IMP.Particle(m)
        p.add_attribute(IMP.FloatKey("thekey"), float(1))
        return p

    def create_particle_pair(self, m):
        p0 = IMP.Particle(m)
        p1 = IMP.Particle(m)
        d0 = IMP.core.XYZ.setup_particle(p0)
        d1 = IMP.core.XYZ.setup_particle(p1)
        d0.set_coordinates(IMP.algebra.Vector3D(0, 0, 1))
        d1.set_coordinates(IMP.algebra.Vector3D(0, 0, 0))
        return (p0, p1)

    def same(self, a, b):
        return a.get_name() == b.get_name()

    def create_singleton_score_state(self, f, a, m, t):
        return IMP.core.SingletonConstraint(f, a, m, t.get_index())

    def create_pair_score_state(self, f, a, m, t):
        return IMP.core.PairConstraint(f, a, m, [x.get_index() for x in t])

    def create_particle_score(self):
        uf = IMP.core.Linear(0, 1)
        return IMP.core.AttributeSingletonScore(uf, IMP.FloatKey("thekey"))

    def create_particle_pair_score(self):
        uf = IMP.core.Linear(0, 1)
        return IMP.core.DistancePairScore(uf)

    def test_set(self):
        """Testing ClassnamesConstraint"""
        # write increment an int field
        # call evaluate and check that it is incremented
        IMP.set_log_level(IMP.VERBOSE)
        print("start")
        m = IMP.Model()
        print("hi")
        c = IMP.container.ListClassnameContainer(m)
        cs = []
        for i in range(0, 30):
            t = self.create_FUNCTIONNAME(m)
            if 'FUNCTIONNAME' == 'particle':
                c.add(t.get_index())
            else:
                c.add([x.get_index() for x in t])
            cs.append(t)
        print("dl")
        k = IMP.IntKey("thevalue")
        f = ClassnameTestModifier(k)  # noqa: F821
        print("apply")
        s = IMP.container.ClassnamesConstraint(f, None, c)
        print("add")
        m.add_score_state(s)
        m.update()
        for p in cs:
            self.assertTrue(FUNCTIONNAME_has_attribute(p, k))  # noqa: F821
        print("done")

    def test_sset(self):
        """Testing ClassnameConstraint"""
        # write increment an int field
        # call evaluate and check that it is incremented
        IMP.set_log_level(IMP.VERBOSE)
        print("start")
        m = IMP.Model()
        print("hi")
        t = self.create_FUNCTIONNAME(m)
        print("dl")
        k = IMP.IntKey("thevalue")
        f = ClassnameTestModifier(k)  # noqa: F821
        print("apply")
        s = self.create_CLASSFUNCTIONNAME_score_state(f, None, m, t)
        m.add_score_state(s)
        print("add")
        m.update()
        self.assertTrue(FUNCTIONNAME_has_attribute(t, k))  # noqa: F821
        print("done")


if __name__ == '__main__':
    IMP.test.main()
