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

TM = coral.getTransactionManager()

L = coral.lang.UML14
DI = coral.lang.XMIDI20

## Checking that original model is OK

prj = coral.loadDocument("testCompositeStateWithMinorStates.xmi")

model = prj.contents()[0]
sm = model.ownedElement.select(lambda x: x.name == "SM")[0]
diagrams = sm.xmidi_get_diagrams()

assert len(diagrams) == 2

d1 = diagrams[0]
d2 = diagrams[1]

assert sm.oclIsKindOf(L.StateMachine)
assert d1.oclIsKindOf(DI.Diagram)
assert d2.oclIsKindOf(DI.Diagram)

assert len(d1.contained) == 2
assert len(d2.contained) == 2

top = sm.top

assert top
assert top.oclIsKindOf(L.CompositeState)

cs1 = top.subvertex.select(lambda x: x.isConcurrent == 0)[0]
cs2 = top.subvertex.select(lambda x: x.isConcurrent == 1)[0]

assert cs1
assert cs2

## The compositestate does not have sub states
assert len(cs1.subvertex) == 0
## The concurrent state has two regions
assert len(cs2.subvertex) == 2

xcs1 = cs1.xmidi_get_graphelements()
xcs2 = cs2.xmidi_get_graphelements()

assert xcs1
assert len(xcs1) == 2
assert xcs2
assert len(xcs2) == 2

ssc1 = xcs1[0].contained.select(lambda x: x.semanticModel.typeInfo == "SubStateCompartment")[0]
ssc2 = xcs1[1].contained.select(lambda x: x.semanticModel.typeInfo == "SubStateCompartment")[0]

assert ssc1
assert ssc2

assert len(ssc1.contained) == 0
assert len(ssc2.contained) == 0

rc1 = xcs2[0].contained.select(lambda x: x.semanticModel.typeInfo == "RegionCompartment")[0]
rc2 = xcs2[1].contained.select(lambda x: x.semanticModel.typeInfo == "RegionCompartment")[0]

assert rc1
assert rc2

assert len(rc1.contained) == 3 #two regions + a separator
assert len(rc1.contained) == 3


## Add a SimpleState into CompositeState
TM.beginModifyCmd()
try:
    cs1.subvertex.insert(L.SimpleState(name="SS1"))
finally:
    TM.endModifyCmd()

assert len(cs1.subvertex) == 1
assert len(ssc1.contained) == 1
assert len(ssc2.contained) == 1

assert len(d1.contained) == 2
assert len(d2.contained) == 2


## Add a SimpleState into CompositeState
TM.beginModifyCmd()
try:
    cs2.subvertex.insert(L.SimpleState(name="SS2"))
finally:
    TM.endModifyCmd()

assert len(cs2.subvertex) == 3 # two old regions + one new
assert len(rc1.contained) == 5 # three regions + two separators
assert len(rc2.contained) == 5

assert len(d1.contained) == 2
assert len(d2.contained) == 2
