#!/usr/bin/env python
import coral
import coral.modeler
import coral.modeler.modeler

TM = coral.getTransactionManager()

## Checking that original model is OK

prj = coral.loadDocument("4states1trans.xml")

class_diagram = prj.contents()[1]
model = prj.contents()[0]

L = coral.lang.UML14
DI = coral.lang.XMIDI20
assert model.oclIsKindOf(L.Model)
assert class_diagram.oclIsKindOf(DI.Diagram)


stm = model.ownedElement.select(lambda x: x.name == "STM")[0]

assert stm.oclIsKindOf(L.StateMachine)
tr = stm.transitions.select(lambda x: x.oclIsKindOf(L.Transition))[0]

g1 = L.Guard(name="G1")
g2 = L.Guard(name="G2")
ce1 = L.CallEvent(name="CE1")
ce2 = L.CallEvent(name="CE2")

TM.beginModifyCmd()
try:
    model.ownedElement.append(ce1)
    model.ownedElement.append(ce2)
finally:
    TM.endModifyCmd()
ca1 = L.CallAction(name="CA1")
ca2 = L.CallAction(name="CA2")


assert tr

graphedges = tr.xmidi_get_graphelements()
assert len(graphedges) == 1

ge = graphedges[0]
td = ge.contained[0]
assert str(td.semanticModel.typeInfo) == "TransitionDescription"

assert len(td.contained) == 0


## Set the effect
print "1"
TM.beginModifyCmd()
try:
    tr.effect = ca1
finally:
    TM.endModifyCmd()

assert tr.effect
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 2

assert str(td.contained[0].semanticModel.typeInfo) == "EffectStart"
assert td.contained[1].semanticModel.element == ca1

## Set the trigger
print "2"
TM.beginModifyCmd()
try:
    tr.trigger = ce1
finally:
    TM.endModifyCmd()

assert tr.effect
assert tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 3

assert td.contained[0].semanticModel.element == ce1
assert str(td.contained[1].semanticModel.typeInfo) == "EffectStart"
assert td.contained[2].semanticModel.element == ca1


## Set the guard
print "3"
TM.beginModifyCmd()
try:
    tr.guard = g1
finally:
    TM.endModifyCmd()

assert tr.effect
assert tr.guard
assert tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 6

assert td.contained[0].semanticModel.element == ce1
assert str(td.contained[1].semanticModel.typeInfo) == "GuardStart"
assert td.contained[2].semanticModel.element == g1
assert str(td.contained[3].semanticModel.typeInfo) == "GuardEnd"
assert str(td.contained[4].semanticModel.typeInfo) == "EffectStart"
assert td.contained[5].semanticModel.element == ca1

## Unset the Guard
print "4"
TM.beginModifyCmd()
try:
    tr.guard = None
finally:
    TM.endModifyCmd()

assert tr.effect
assert not tr.guard
assert tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 3

assert td.contained[0].semanticModel.element == ce1
assert str(td.contained[1].semanticModel.typeInfo) == "EffectStart"
assert td.contained[2].semanticModel.element == ca1

## Unset Trigger
print "5"
TM.beginModifyCmd()
try:
    tr.trigger = None
finally:
    TM.endModifyCmd()

assert tr.effect
assert not tr.guard
assert not tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 2

assert str(td.contained[0].semanticModel.typeInfo) == "EffectStart"
assert td.contained[1].semanticModel.element == ca1


## Put the Guard back (new element)
print "6"
TM.beginModifyCmd()
try:
    tr.guard = g2
finally:
    TM.endModifyCmd()

assert tr.effect
assert tr.guard
assert not tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 5

assert str(td.contained[0].semanticModel.typeInfo) == "GuardStart"
assert td.contained[1].semanticModel.element == g2
assert str(td.contained[2].semanticModel.typeInfo) == "GuardEnd"
assert str(td.contained[3].semanticModel.typeInfo) == "EffectStart"
assert td.contained[4].semanticModel.element == ca1


## Unset effect
print "7"
TM.beginModifyCmd()
try:
    tr.effect = None
finally:
    TM.endModifyCmd()

assert not tr.effect
assert tr.guard
assert not tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 3

assert str(td.contained[0].semanticModel.typeInfo) == "GuardStart"
assert td.contained[1].semanticModel.element == g2
assert str(td.contained[2].semanticModel.typeInfo) == "GuardEnd"

## Unset guard, set effect and trigger
print "8"
TM.beginModifyCmd()
try:
    tr.effect = ca2
    tr.trigger = ce2
    tr.guard = None
finally:
    TM.endModifyCmd()

assert tr.effect
assert not tr.guard
assert tr.trigger
for i in td.contained:
    print i, i.getName()
assert len(td.contained) == 3

assert td.contained[0].semanticModel.element == ce2
assert str(td.contained[1].semanticModel.typeInfo) == "EffectStart"
assert td.contained[2].semanticModel.element == ca2
