From d35a165e63cc47f0f8ce5194df6cd12286dbc2dd Mon Sep 17 00:00:00 2001
From: Edwin Lai <edwinslai@gmail.com>
Date: Wed, 2 Aug 2023 09:34:06 -0400
Subject: [PATCH] fix: preserve stratify drilldown state when navigating
 forward/backward  (#1651)

---
 .../components/models/tera-model-diagram.vue  |  26 ++-
 .../models/tera-model-type-legend.vue         |  86 ++++++++
 .../models/tera-reflexives-toolbar.vue        | 187 ++++++++++--------
 .../models/tera-strata-model-diagram.vue      |  76 +------
 .../models/tera-typed-model-diagram.vue       | 184 +++++++++--------
 .../src/components/workflow/tera-stratify.vue | 111 +++++++----
 .../petrinet/petrinet-service.ts              |   6 -
 .../services/models/stratification-service.ts |  18 +-
 .../services/stratification-service.spec.ts   | 109 +++++-----
 9 files changed, 443 insertions(+), 360 deletions(-)
 create mode 100644 packages/client/hmi-client/src/components/models/tera-model-type-legend.vue

diff --git a/packages/client/hmi-client/src/components/models/tera-model-diagram.vue b/packages/client/hmi-client/src/components/models/tera-model-diagram.vue
index 178cde353a..6aa75246f5 100644
--- a/packages/client/hmi-client/src/components/models/tera-model-diagram.vue
+++ b/packages/client/hmi-client/src/components/models/tera-model-diagram.vue
@@ -45,8 +45,15 @@
 										"
 									/>
 								</span>
+								<Button
+									v-if="model && isStratifiedAMR(model) && !isEditing"
+									@click="toggleCollapsedView"
+									:label="isCollapsed ? 'Show expanded view' : 'Show collapsed view'"
+									class="p-button-sm p-button-outlined toolbar-button"
+								/>
 							</template>
 						</Toolbar>
+						<tera-model-type-legend v-if="model" :model="model" />
 						<div v-if="model" ref="graphElement" class="graph-element" />
 						<ContextMenu ref="menu" :model="contextMenuItems" />
 					</section>
@@ -226,6 +233,7 @@ import { Model, Observable } from '@/types/Types';
 import TeraModal from '@/components/widgets/tera-modal.vue';
 import InputText from 'primevue/inputtext';
 import TeraResizablePanel from '../widgets/tera-resizable-panel.vue';
+import TeraModelTypeLegend from './tera-model-type-legend.vue';
 
 // Get rid of these emits
 const emit = defineEmits([
@@ -457,8 +465,22 @@ const contextMenuItems = ref([
 	}
 ]);
 
+const isCollapsed = ref(true);
+async function toggleCollapsedView() {
+	isCollapsed.value = !isCollapsed.value;
+	if (props.model) {
+		const graphData: IGraph<NodeData, EdgeData> = convertToIGraphHelper(props.model);
+		// Render graph
+		if (renderer) {
+			renderer.isGraphDirty = true;
+			await renderer.setData(graphData);
+			await renderer.render();
+		}
+	}
+}
+
 const convertToIGraphHelper = (amr: Model) => {
-	if (isStratifiedAMR(amr)) {
+	if (isStratifiedAMR(amr) && isCollapsed.value) {
 		// FIXME: wont' work for MIRA
 		return convertToIGraph(props.model?.semantics?.span?.[0].system);
 	}
@@ -592,7 +614,7 @@ const cancelEdit = async () => {
 	if (!props.model) return;
 
 	// Convert petri net into a graph with raw input data
-	const graphData: IGraph<NodeData, EdgeData> = convertToIGraph(props.model);
+	const graphData: IGraph<NodeData, EdgeData> = convertToIGraphHelper(props.model);
 
 	if (renderer) {
 		renderer.setEditMode(false);
diff --git a/packages/client/hmi-client/src/components/models/tera-model-type-legend.vue b/packages/client/hmi-client/src/components/models/tera-model-type-legend.vue
new file mode 100644
index 0000000000..9fc951e5bd
--- /dev/null
+++ b/packages/client/hmi-client/src/components/models/tera-model-type-legend.vue
@@ -0,0 +1,86 @@
+<template>
+	<section class="legend" v-if="stateTypes || transitionTypes">
+		<ul>
+			<li v-for="(type, i) in stateTypes" :key="i">
+				<div class="legend-key-circle" :style="getLegendKeyStyle(type)" />
+				{{ type }}
+			</li>
+		</ul>
+		<ul>
+			<li v-for="(type, i) in transitionTypes" :key="i">
+				<div class="legend-key-square" :style="getLegendKeyStyle(type)" />
+				{{ type }}
+			</li>
+		</ul>
+	</section>
+</template>
+
+<script setup lang="ts">
+import { computed } from 'vue';
+import { Model } from '@/types/Types';
+import { useNodeTypeColorPalette } from '@/utils/petrinet-color-palette';
+
+const props = defineProps<{
+	model: Model;
+}>();
+
+const { getNodeTypeColor } = useNodeTypeColorPalette();
+
+const stateTypes = computed<string[]>(() =>
+	props.model.semantics?.typing?.system?.model.states.map((s) => s.name)
+);
+const transitionTypes = computed<string[]>(() =>
+	props.model.semantics?.typing?.system?.model.transitions.map((t) => t.properties?.name)
+);
+
+function getLegendKeyStyle(id: string) {
+	if (!id) {
+		return {
+			backgroundColor: 'var(--petri-nodeFill)'
+		};
+	}
+	return {
+		backgroundColor: getNodeTypeColor(id)
+	};
+}
+</script>
+
+<style scoped>
+.legend {
+	position: absolute;
+	bottom: 0;
+	z-index: 1;
+	margin-bottom: 1rem;
+	margin-left: 1rem;
+	display: flex;
+	gap: 1rem;
+	background-color: var(--surface-section);
+	border-radius: 0.5rem;
+	padding: 0.5rem;
+}
+
+.legend-key-circle {
+	height: 24px;
+	width: 24px;
+	border-radius: 12px;
+}
+
+.legend-key-square {
+	height: 24px;
+	width: 24px;
+	border-radius: 4px;
+}
+
+section.legend ul {
+	display: flex;
+	gap: 0.5rem;
+	list-style-type: none;
+	flex-direction: row;
+}
+
+section.legend li {
+	display: flex;
+	align-items: center;
+	gap: 0.5rem;
+}
+</style>
diff --git a/packages/client/hmi-client/src/components/models/tera-reflexives-toolbar.vue b/packages/client/hmi-client/src/components/models/tera-reflexives-toolbar.vue
index bcb05b4323..72288c1d07 100644
--- a/packages/client/hmi-client/src/components/models/tera-reflexives-toolbar.vue
+++ b/packages/client/hmi-client/src/components/models/tera-reflexives-toolbar.vue
@@ -12,7 +12,7 @@
 					optionLabel="id"
 					:model-value="statesToAddReflexives[transition.id]"
 					@update:model-value="
-						(newValue) => updateStatesToAddReflexives(newValue, transition, stateType as string)
+						(states) => updateStatesToAddReflexives({states, typeOfTransition: transition, typeIdOfState: stateType as string}, i)
 					"
 				/>
 			</div>
@@ -29,6 +29,7 @@ import {
 	addTyping,
 	updateRateExpression
 } from '@/model-representation/petrinet/petrinet-service';
+import { cloneDeep } from 'lodash';
 
 const props = defineProps<{
 	modelToUpdate: Model; // the model to which we will add reflexives
@@ -40,8 +41,7 @@ const emit = defineEmits(['model-updated']);
 const modelToCompareTypeSystem = computed<TypeSystem | undefined>(
 	() => props.modelToCompare.semantics?.typing?.system.model
 );
-const typedModel = ref<Model>(props.modelToUpdate); // this is the object that is being edited
-let unassignedTransitionTypes: Transition[] = [];
+const typedModel = ref<Model>(cloneDeep(props.modelToUpdate)); // this is the object that is being edited
 const statesToAddReflexives = ref<{ [id: string]: { id: string; name: string }[] }>({});
 const typeIdToTransitionIdMap = computed<{ [id: string]: string }>(() => {
 	const map: { [id: string]: string } = {};
@@ -63,7 +63,6 @@ const reflexiveNodeOptions = computed<{ [id: string]: { id: string; name: string
 	});
 	return options;
 });
-
 const stateId2NameMap = computed<{ [id: string]: string }>(() => {
 	const map: { [id: string]: string } = {};
 	props.modelToUpdate.model.states.forEach((state) => {
@@ -72,96 +71,118 @@ const stateId2NameMap = computed<{ [id: string]: string }>(() => {
 	return map;
 });
 
+const addedReflexivesRows: {
+	states: {
+		id: string;
+		name: string;
+	}[];
+	typeOfTransition: Transition;
+	typeIdOfState: string;
+}[] = [];
+
+/* Every time the user changes their selection of what states to add reflexives to, overwrite the previous changes with the current ones.
+   Since there can be multiple MultiSelect components, the selections for all MultiSelect components are combined in 'addedReflexivesRows'.
+   Iterate through 'addedReflexivesRows' and add reflexives according to the selections. 
+*/
 function updateStatesToAddReflexives(
-	newValue: {
-		id: string; // id of the state to which to add reflexive
-		name: string; // name of the state to which to add reflexive
-	}[],
-	typeOfTransition: Transition, // e.g. infect, recover
-	typeIdOfState: string // e.g. pop
+	selection: {
+		states: {
+			id: string;
+			name: string;
+		}[]; // list of id+name of state to which to add reflexives
+		typeOfTransition: Transition; // e.g. infect, recover
+		typeIdOfState: string; // e.g. pop
+	},
+	index: number
 ) {
-	statesToAddReflexives.value[typeOfTransition.id] = newValue;
-	const updatedTypeMap = typedModel.value.semantics?.typing?.map;
-	const updatedTypeSystem = typedModel.value.semantics?.typing?.system;
-
-	if (updatedTypeMap && updatedTypeSystem) {
-		newValue.forEach((state) => {
-			const newTransitionId = `${typeIdToTransitionIdMap.value[typeOfTransition.id]}${state.id}${
-				state.id
-			}`;
-			// For the type of reflexive transition that we are adding, get the number of inputs and outputs that share the same type as the state that we are updating
-			// E.g. if we are adding an 'Infect' reflexive to a node of type 'Pop', get the number of 'Pop' inputs and outputs for 'Infect'
-			const numInputsOfStateType = typeOfTransition.input.filter((i) => i === typeIdOfState).length;
-			// const numOutputsOfStateType = typeOfTransition.input.filter(i => i === typeIdOfState).length;
-			// Assume for now that the number of inputs and outputs for a given type are always equal, though in general this may not be the case
-			// TODO: implement logic for more generalized case where the above assumption is not true
-			addReflexives(typedModel.value, state.id, newTransitionId, numInputsOfStateType);
-			const reflexive = typedModel.value.model.transitions.find((t) => t.id === newTransitionId);
-
-			const transition = props.modelToCompare?.semantics?.typing?.system.model.transitions.find(
-				(t) => t.id === typeOfTransition.id
-			);
-			if (transition) {
-				updateRateExpression(typedModel.value, reflexive as PetriNetTransition, '');
-				if (!updatedTypeMap.find((m) => m[0] === newTransitionId)) {
-					updatedTypeMap.push([newTransitionId, typeOfTransition.id]);
+	typedModel.value = cloneDeep(props.modelToUpdate);
+	addedReflexivesRows[index] = selection;
+	addedReflexivesRows.forEach(({ states, typeOfTransition, typeIdOfState }) => {
+		statesToAddReflexives.value[typeOfTransition.id] = states;
+		const updatedTypeMap = typedModel.value.semantics?.typing?.map;
+		const updatedTypeSystem = typedModel.value.semantics?.typing?.system;
+		if (updatedTypeMap && updatedTypeSystem) {
+			states.forEach((state) => {
+				// For the type of reflexive transition that we are adding, get the number of inputs and outputs that share the same type as the state that we are updating
+				// E.g. if we are adding an 'Infect' reflexive to a node of type 'Pop', get the number of 'Pop' inputs and outputs for 'Infect'
+				const newTransitionId = `${typeIdToTransitionIdMap.value[typeOfTransition.id]}${state.id}`;
+
+				// Assume for now that the number of inputs and outputs for a given type are always equal, though in general this may not be the case
+				// TODO: implement logic for more generalized case where the above assumption is not true
+				const numInputsOfStateType = typeOfTransition.input.filter(
+					(i) => i === typeIdOfState
+				).length;
+				// const numOutputsOfStateType = typeOfTransition.input.filter(i => i === typeIdOfState).length;
+
+				if (!typedModel.value.model.transitions.find((t) => t.id === newTransitionId)) {
+					addReflexives(typedModel.value, state.id, newTransitionId, numInputsOfStateType);
 				}
-				if (!updatedTypeSystem.model.transitions.find((t) => t.id === typeOfTransition.id)) {
-					updatedTypeSystem.model.transitions.push(transition);
+				const reflexive = typedModel.value.model.transitions.find((t) => t.id === newTransitionId);
+
+				const transition = props.modelToCompare?.semantics?.typing?.system.model.transitions.find(
+					(t) => t.id === typeOfTransition.id
+				);
+				if (transition) {
+					updateRateExpression(typedModel.value, reflexive as PetriNetTransition, '');
+					if (!updatedTypeMap.find((m) => m[0] === newTransitionId)) {
+						updatedTypeMap.push([newTransitionId, typeOfTransition.id]);
+					}
+					if (!updatedTypeSystem.model.transitions.find((t) => t.id === typeOfTransition.id)) {
+						updatedTypeSystem.model.transitions.push(transition);
+					}
 				}
-			}
-		});
-		const updatedTyping: TypingSemantics = {
-			map: updatedTypeMap,
-			system: updatedTypeSystem
-		};
-		addTyping(typedModel.value, updatedTyping);
-	}
+			});
+			const updatedTyping: TypingSemantics = {
+				map: updatedTypeMap,
+				system: updatedTypeSystem
+			};
+			addTyping(typedModel.value, updatedTyping);
+		}
+	});
 	emit('model-updated', typedModel.value);
 }
 
-watch(
-	() => props.modelToUpdate,
-	() => {
-		if (props.modelToCompare) {
-			typedModel.value = props.modelToUpdate;
-			emit('model-updated', typedModel.value);
+/* Compare the type systems of 'modelToUpdate' and 'modelToCompare' to determine what options the user has for adding reflexives.
+   Allow the user to add transitions types that are present in 'modelToCompare' but not in 'modelToUpdate'.
+   E.g. if 'modelToUpdate' has ['Pop','Infect'] transitions and 'modelToCompare' has ['Pop,'Infect','Recover'] transitions, 
+   user should be prompted to add 'Recover' transitions to 'modelToUpdate'
+*/
+function populateReflexiveOptions() {
+	if (modelToCompareTypeSystem.value) {
+		let unassignedTransitions: Transition[];
+		const modelToUpdateTransitionIds =
+			props.modelToUpdate.semantics?.typing?.system.model.transitions.map((t) => t.id);
+		const modelToCompareTypeTransitionIds = modelToCompareTypeSystem.value?.transitions.map(
+			(t) => t.id
+		);
+		if (modelToUpdateTransitionIds && modelToCompareTypeTransitionIds) {
+			const unassignedIds = modelToCompareTypeTransitionIds.filter(
+				(id) => !modelToUpdateTransitionIds.includes(id)
+			);
+			// get the transition types that are in 'modelToCompare' but not 'modelToUpdate'
+			unassignedTransitions = modelToCompareTypeSystem.value?.transitions.filter((t) =>
+				unassignedIds.includes(t.id)
+			);
 		}
-	},
-	{ immediate: true }
-);
+		props.modelToUpdate.model.states.forEach((state) => {
+			// get type of state for each state in model to update model
+			const type: string =
+				props.modelToUpdate.semantics?.typing?.map.find((m) => m[0] === state.id)?.[1] ?? '';
+			// for each unassigned transition type, check if inputs or ouputs have the type of this state
+			// you should only be allowed to add a transition to a state, if the transition has inputs or outputs of the same type as the state
+			const allowedTransitionsForState: Transition[] = unassignedTransitions.filter(
+				(unassigned) => unassigned.input.includes(type) || unassigned.output.includes(type)
+			);
+			reflexiveOptions.value[type] = allowedTransitionsForState;
+		});
+	}
+}
 
 watch(
-	[() => modelToCompareTypeSystem],
+	[() => props.modelToCompare, () => props.modelToUpdate.semantics?.typing],
 	() => {
-		if (modelToCompareTypeSystem.value) {
-			const modelToUpdateTransitionIds =
-				props.modelToUpdate.semantics?.typing?.system.model.transitions.map((t) => t.id);
-			const modelToCompareTypeTransitionIds = modelToCompareTypeSystem.value?.transitions.map(
-				(t) => t.id
-			);
-			if (modelToUpdateTransitionIds && modelToCompareTypeTransitionIds) {
-				const unassignedIds = modelToCompareTypeTransitionIds.filter(
-					(id) => !modelToUpdateTransitionIds.includes(id)
-				);
-				const unassignedTransitions: Transition[] =
-					modelToCompareTypeSystem.value?.transitions.filter((t) => unassignedIds.includes(t.id));
-				if (unassignedTransitions.length > 0) {
-					unassignedTransitionTypes = unassignedTransitionTypes.concat(unassignedTransitions);
-				}
-			}
-			props.modelToUpdate.model.states.forEach((state) => {
-				// get type of state for each state in model to update model
-				const type: string =
-					props.modelToUpdate.semantics?.typing?.map.find((m) => m[0] === state.id)?.[1] ?? '';
-				// for each unassigned transition type, check if inputs or ouputs have the type of this state
-				const allowedTransitionsForState: Transition[] = unassignedTransitionTypes.filter(
-					(unassigned) => unassigned.input.includes(type) || unassigned.output.includes(type)
-				);
-				if (!reflexiveOptions.value[type]) {
-					reflexiveOptions.value[type] = allowedTransitionsForState;
-				}
-			});
+		if (props.modelToUpdate.semantics?.typing) {
+			populateReflexiveOptions();
 		}
 	},
 	{ immediate: true }
diff --git a/packages/client/hmi-client/src/components/models/tera-strata-model-diagram.vue b/packages/client/hmi-client/src/components/models/tera-strata-model-diagram.vue
index c87a33d8ea..c70942f73b 100644
--- a/packages/client/hmi-client/src/components/models/tera-strata-model-diagram.vue
+++ b/packages/client/hmi-client/src/components/models/tera-strata-model-diagram.vue
@@ -4,30 +4,18 @@
 			<div ref="splitterContainer" class="splitter-container">
 				<section class="graph-element">
 					<tera-reflexives-toolbar
-						v-if="showReflexivesToolbar && strataModel && baseModel"
+						v-if="strataModel && baseModel"
+						v-show="showReflexivesToolbar"
 						:model-to-update="strataModel"
 						:model-to-compare="baseModel"
 						@model-updated="
 							(value) => {
 								typedModel = value;
-								emit('model-updated', value);
+								emit('model-updated', typedModel);
 							}
 						"
 					/>
-					<section class="legend">
-						<ul>
-							<li v-for="(type, i) in stateTypes" :key="i">
-								<div class="legend-key-circle" :style="getLegendKeyStyle(type ?? '')" />
-								{{ type }}
-							</li>
-						</ul>
-						<ul>
-							<li v-for="(type, i) in transitionTypes" :key="i">
-								<div class="legend-key-square" :style="getLegendKeyStyle(type ?? '')" />
-								{{ type }}
-							</li>
-						</ul>
-					</section>
+					<tera-model-type-legend :model="typedModel" />
 					<div v-if="typedModel" ref="graphElement" class="graph-element" />
 				</section>
 			</div>
@@ -37,7 +25,7 @@
 
 <script setup lang="ts">
 import { IGraph } from '@graph-scaffolder/index';
-import { watch, ref, computed } from 'vue';
+import { watch, ref } from 'vue';
 import { runDagreLayout } from '@/services/graph';
 import {
 	PetrinetRenderer,
@@ -49,10 +37,11 @@ import { Model, TypeSystem } from '@/types/Types';
 import { useNodeTypeColorPalette } from '@/utils/petrinet-color-palette';
 import TeraResizablePanel from '../widgets/tera-resizable-panel.vue';
 import TeraReflexivesToolbar from './tera-reflexives-toolbar.vue';
+import TeraModelTypeLegend from './tera-model-type-legend.vue';
 
 const props = defineProps<{
 	strataModel: Model;
-	baseModel?: Model;
+	baseModel?: Model | null;
 	baseModelTypeSystem?: TypeSystem;
 	showReflexivesToolbar: boolean;
 }>();
@@ -66,20 +55,7 @@ const splitterContainer = ref<HTMLElement | null>(null);
 const graphElement = ref<HTMLDivElement | null>(null);
 let renderer: PetrinetRenderer | null = null;
 
-const stateTypes = computed(() =>
-	typedModel.value.semantics?.typing?.system?.model.states.map((s) => s.name)
-);
-const transitionTypes = computed(() =>
-	typedModel.value.semantics?.typing?.system?.model.transitions.map((t) => t.properties?.name)
-);
-
-const { getNodeTypeColor, setNodeTypeColor } = useNodeTypeColorPalette();
-
-function getLegendKeyStyle(id: string) {
-	return {
-		backgroundColor: getNodeTypeColor(id)
-	};
-}
+const { setNodeTypeColor } = useNodeTypeColorPalette();
 
 // Whenever selectedModelId changes, fetch model with that ID
 watch(
@@ -140,42 +116,6 @@ main {
 	overflow: auto;
 }
 
-.legend {
-	position: absolute;
-	bottom: 0;
-	z-index: 1;
-	margin-bottom: 1rem;
-	margin-left: 1rem;
-	display: flex;
-	gap: 1rem;
-	background-color: var(--surface-section);
-	border-radius: 0.5rem;
-	padding: 0.5rem;
-}
-
-.legend-key-circle {
-	height: 24px;
-	width: 24px;
-	border-radius: 12px;
-}
-
-.legend-key-square {
-	height: 24px;
-	width: 24px;
-	border-radius: 4px;
-}
-
-ul {
-	display: flex;
-	gap: 0.5rem;
-	list-style-type: none;
-}
-
-li {
-	display: flex;
-	align-items: center;
-	gap: 0.5rem;
-}
 .splitter-container {
 	height: 100%;
 }
diff --git a/packages/client/hmi-client/src/components/models/tera-typed-model-diagram.vue b/packages/client/hmi-client/src/components/models/tera-typed-model-diagram.vue
index ba3cd9f328..884435b9cc 100644
--- a/packages/client/hmi-client/src/components/models/tera-typed-model-diagram.vue
+++ b/packages/client/hmi-client/src/components/models/tera-typed-model-diagram.vue
@@ -48,7 +48,7 @@
 							</div>
 							<!-- cancel row  -->
 							<div>
-								<Button icon="pi pi-times" text rounded />
+								<Button icon="pi pi-times" text rounded @click="cancelTypedRow(index)" />
 							</div>
 						</div>
 						<Button
@@ -60,7 +60,8 @@
 						/>
 					</section>
 					<tera-reflexives-toolbar
-						v-if="showReflexivesToolbar && model && strataModel"
+						v-if="model && strataModel"
+						v-show="showReflexivesToolbar"
 						:model-to-update="model"
 						:model-to-compare="strataModel"
 						@model-updated="
@@ -70,20 +71,7 @@
 							}
 						"
 					/>
-					<section class="legend">
-						<ul>
-							<li v-for="(type, i) in stateTypes" :key="i">
-								<div class="legend-key-circle" :style="getLegendKeyStyle(type ?? '')" />
-								{{ type }}
-							</li>
-						</ul>
-						<ul>
-							<li v-for="(type, i) in transitionTypes" :key="i">
-								<div class="legend-key-square" :style="getLegendKeyStyle(type ?? '')" />
-								{{ type }}
-							</li>
-						</ul>
-					</section>
+					<tera-model-type-legend :model="typedModel" />
 					<Toolbar>
 						<template #end>
 							<Button
@@ -129,8 +117,9 @@ import { NestedPetrinetRenderer } from '@/model-representation/petrinet/nested-p
 import Toolbar from 'primevue/toolbar';
 import TeraResizablePanel from '../widgets/tera-resizable-panel.vue';
 import TeraReflexivesToolbar from './tera-reflexives-toolbar.vue';
+import TeraModelTypeLegend from './tera-model-type-legend.vue';
 
-const emit = defineEmits(['model-updated', 'all-nodes-typed']);
+const emit = defineEmits(['model-updated', 'all-nodes-typed', 'not-all-nodes-typed']);
 
 const props = defineProps<{
 	model: Model;
@@ -144,12 +133,6 @@ const graphElement = ref<HTMLDivElement | null>(null);
 let renderer: PetrinetRenderer | null = null;
 
 const typedModel = ref<Model>(props.model);
-const stateTypes = computed(() =>
-	props.model.semantics?.typing?.system?.model.states.map((s) => s.name)
-);
-const transitionTypes = computed(() =>
-	props.model.semantics?.typing?.system?.model.transitions.map((t) => t.properties?.name)
-);
 // these are values that user will edit/select that correspond to each row in the model typing editor
 const typedRows = ref<
 	{
@@ -173,10 +156,14 @@ const assignToOptions = computed<{ [s: string]: string[] }[]>(() => {
 	});
 	return options;
 });
+const allNodesTyped = computed(() => numberTypedNodes.value === numberNodes.value);
 
 function addTypedRow() {
 	typedRows.value.push({});
 }
+function cancelTypedRow(rowIndex: number) {
+	typedRows.value.splice(rowIndex, 1);
+}
 // 'typedRows.typeName' is assigned the value of 'typeNameBuffer' when the user focuses away from the associated InputText
 let typeNameBuffer: string[] = [];
 function setTypeNameBuffer(newValue, row) {
@@ -248,23 +235,61 @@ watch(
 );
 
 watch(
-	() => props.typeSystem,
+	() => [props.typeSystem, props.showTypingToolbar],
 	() => {
-		setNodeColors();
-		if (typedRows.value.length === 0) {
-			typedRows.value.push(
-				{
-					nodeType: 'Variable',
-					typeName: props.typeSystem?.states[0].name
-				},
-				{
-					nodeType: 'Transition',
-					typeName: props.typeSystem?.transitions[0].properties?.name
-				}
-			);
+		if (props.showTypingToolbar) {
+			setNodeColors();
+			if (typedModel.value.semantics?.typing) {
+				// pre-populate 'typedRows' if 'typedModel' already has typing
+				const typedRowsToPopulate: {
+					nodeType: string;
+					typeName: string;
+					assignTo: string[];
+				}[] = [];
+				const typedRowsTypeNames: Set<string> = new Set();
+				typedModel.value.semantics.typing.map.forEach((m) => {
+					const nodeId = m[0];
+					const typeId = m[1];
+					const state = typedModel.value.model.states.find((s) => s.id === nodeId);
+					const transition = typedModel.value.model.transitions.find((t) => t.id === nodeId);
+					const typeState = typedModel.value.semantics!.typing!.system.model.states.find(
+						(s) => s.id === typeId
+					);
+					const typeTransition = typedModel.value.semantics!.typing!.system.model.transitions.find(
+						(t) => t.id === typeId
+					);
+					const node = state || transition;
+					const nodeType = state ? 'Variable' : 'Transition';
+					const typeName = typeState ? typeState.id : typeTransition.properties?.name;
+					if (!typedRowsTypeNames.has(typeName)) {
+						typedRowsTypeNames.add(typeName);
+						typedRowsToPopulate.push({
+							nodeType,
+							typeName,
+							assignTo: [node.id]
+						});
+					} else {
+						const assignTo = typedRowsToPopulate.find((row) => row.typeName === typeName)?.assignTo;
+						assignTo?.push(node.id);
+					}
+				});
+				typedRows.value = typedRowsToPopulate;
+			} else if (typedRows.value.length === 0) {
+				typedRows.value.push(
+					{
+						nodeType: 'Variable',
+						typeName: props.typeSystem?.states[0].name
+					},
+					{
+						nodeType: 'Transition',
+						typeName: props.typeSystem?.transitions[0].properties?.name
+					}
+				);
+			}
+			typeNameBuffer = typedRows.value.map((r) => r.typeName ?? '');
 		}
-		typeNameBuffer = typedRows.value.map((r) => r.typeName ?? '');
-	}
+	},
+	{ immediate: true }
 );
 
 // construct TypingSemantics data structure when user updates variable/transition assignments
@@ -307,21 +332,19 @@ watch(
 			}
 		});
 
-		if (stateTypedMap.length > 0) {
-			// eslint-disable-next-line @typescript-eslint/naming-convention
-			const { name, description, schema, model_version } = typedModel.value;
-			typingSemantics = {
-				map: stateTypedMap,
-				system: {
-					name,
-					description,
-					schema,
-					model_version,
-					model: updatedTypeSystem
-				}
-			};
-			addTyping(typedModel.value, typingSemantics);
-		}
+		// eslint-disable-next-line @typescript-eslint/naming-convention
+		const { name, description, schema, model_version } = typedModel.value;
+		typingSemantics = {
+			map: stateTypedMap,
+			system: {
+				name,
+				description,
+				schema,
+				model_version,
+				model: updatedTypeSystem
+			}
+		};
+		addTyping(typedModel.value, typingSemantics);
 
 		transitionTypedMap.forEach((map) => {
 			// See if there is a corresponding type defined in the strata model's type system
@@ -343,7 +366,6 @@ watch(
 		});
 		if (transitionTypedMap.length > 0) {
 			// eslint-disable-next-line @typescript-eslint/naming-convention
-			const { name, description, schema, model_version } = typedModel.value;
 			const typeMap: string[][] = [...stateTypedMap, ...transitionTypedMap];
 			typingSemantics = {
 				map: typeMap,
@@ -364,8 +386,10 @@ watch(
 watch(
 	numberTypedNodes,
 	() => {
-		if (numberTypedNodes.value === numberNodes.value) {
+		if (allNodesTyped.value) {
 			emit('all-nodes-typed', typedModel.value);
+		} else {
+			emit('not-all-nodes-typed', numberNodes.value - numberTypedNodes.value);
 		}
 	},
 	{ immediate: true }
@@ -431,44 +455,6 @@ main {
 	border-radius: var(--border-radius);
 	overflow: auto;
 }
-
-.legend {
-	position: absolute;
-	bottom: 0;
-	z-index: 1;
-	margin-bottom: 1rem;
-	margin-left: 1rem;
-	display: flex;
-	gap: 1rem;
-	background-color: var(--surface-section);
-	border-radius: 0.5rem;
-	padding: 0.5rem;
-}
-
-.legend-key-circle {
-	height: 24px;
-	width: 24px;
-	border-radius: 12px;
-}
-
-.legend-key-square {
-	height: 24px;
-	width: 24px;
-	border-radius: 4px;
-}
-
-ul {
-	display: flex;
-	gap: 0.5rem;
-	list-style-type: none;
-}
-
-li {
-	display: flex;
-	align-items: center;
-	gap: 0.5rem;
-}
-
 .splitter-container {
 	height: 100%;
 }
@@ -549,4 +535,16 @@ li {
 	background-color: var(--surface-0);
 	margin: 0.25rem;
 }
+
+.legend-key-circle {
+	height: 24px;
+	width: 24px;
+	border-radius: 12px;
+}
+
+.legend-key-square {
+	height: 24px;
+	width: 24px;
+	border-radius: 4px;
+}
 </style>
diff --git a/packages/client/hmi-client/src/components/workflow/tera-stratify.vue b/packages/client/hmi-client/src/components/workflow/tera-stratify.vue
index 264db04863..ddbe93db7d 100644
--- a/packages/client/hmi-client/src/components/workflow/tera-stratify.vue
+++ b/packages/client/hmi-client/src/components/workflow/tera-stratify.vue
@@ -16,17 +16,40 @@
 					icon="pi pi-sign-out"
 					@click="stratifyView = StratifyView.Output"
 					:active="stratifyView === StratifyView.Output"
+					:disabled="!stratifiedModel"
 				/>
 			</span>
-			<Button
-				v-if="stratifyView === StratifyView.Input"
-				class="stratify-button p-button-sm"
-				label="Stratify"
-				icon="pi pi-arrow-right"
-				iconPos="right"
-				@click="doStratify"
-				:disabled="stratifyStep !== 3"
-			/>
+			<div class="buttons" v-if="strataModel && stratifyView === StratifyView.Input">
+				<Button
+					class="p-button-outlined"
+					label="Go back"
+					icon="pi pi-arrow-left"
+					:disabled="stratifyStep === 0"
+					@click="goBack"
+				/>
+				<Button
+					v-if="stratifyStep === 1"
+					label="Continue to step 2: Assign types"
+					icon="pi pi-arrow-right"
+					iconPos="right"
+					@click="stratifyStep = 2"
+				/>
+				<Button
+					v-if="typedBaseModel && stratifyStep === 2"
+					label="Continue to step 3: Manage interactions"
+					icon="pi pi-arrow-right"
+					iconPos="right"
+					@click="stratifyStep = 3"
+				/>
+				<Button
+					v-if="stratifyStep === 3"
+					class="stratify-button p-button-sm"
+					label="Stratify"
+					icon="pi pi-arrow-right"
+					iconPos="right"
+					@click="doStratify"
+				/>
+			</div>
 		</header>
 		<section v-if="stratifyView === StratifyView.Input">
 			<nav>
@@ -45,30 +68,20 @@
 			</nav>
 			<section class="step-1">
 				<div class="instructions">
-					<div class="buttons" v-if="strataModel">
-						<Button
-							class="p-button-outlined"
-							label="Go back"
-							icon="pi pi-arrow-left"
-							:disabled="stratifyStep === 0"
-							@click="goBack"
-						/>
-						<Button
-							v-if="stratifyStep === 1"
-							label="Continue to step 2: Assign types"
-							icon="pi pi-arrow-right"
-							iconPos="right"
-							@click="stratifyStep = 2"
-						/>
-						<Button
-							v-if="typedBaseModel && stratifyStep === 2"
-							label="Continue to step 3: Manage interactions"
-							icon="pi pi-arrow-right"
-							iconPos="right"
-							@click="stratifyStep = 3"
-						/>
-					</div>
-					<span v-else>Define the groups you want to stratify your model with.</span>
+					<span v-if="stratifyStep === 1"
+						>Define the groups you want to stratify your model with.</span
+					>
+					<span v-if="stratifyStep === 2 && numUntypedNodes > 0"
+						>Assign types to each of the nodes in your model.
+					</span>
+					<span v-if="stratifyStep === 2 && numUntypedNodes > 0" class="nodes-require-types"
+						>{{ numUntypedNodes }} nodes require types</span
+					>
+					<span v-if="stratifyStep === 2 && numUntypedNodes === 0">All nodes have types.</span>
+					<span v-if="stratifyStep === 3"
+						>Make sure all behaviour in the original model are in the stratified model in order for
+						them to be part of the final model.</span
+					>
 				</div>
 				<div class="step-1-inner">
 					<Accordion :active-index="0">
@@ -81,6 +94,12 @@
 								:type-system="strataModel?.semantics?.typing?.system.model"
 								@model-updated="(value) => (typedBaseModel = value)"
 								@all-nodes-typed="(value) => onAllNodesTyped(value)"
+								@not-all-nodes-typed="
+									(value) => {
+										typedBaseModel = null;
+										numUntypedNodes = value;
+									}
+								"
 								:show-reflexives-toolbar="stratifyStep === 3"
 							/>
 						</AccordionTab>
@@ -227,13 +246,14 @@ const labels = ref();
 const strataModel = ref<Model | null>(null);
 const modelConfiguration = ref<ModelConfiguration>();
 const model = ref<Model | null>(null);
-const typedBaseModel = ref<Model>();
+const typedBaseModel = ref<Model | null>(null);
 const typedStrataModel = ref<Model | null>(null);
 const stratifiedModel = ref<Model>();
+const numUntypedNodes = ref();
 
 function onAllNodesTyped(value: Model) {
-	console.log(value);
 	typedBaseModel.value = value;
+	numUntypedNodes.value = 0;
 }
 
 function generateStrataModel() {
@@ -273,11 +293,10 @@ async function doStratify() {
 }
 
 function goBack() {
-	if (stratifyStep.value > 0) {
-		if (stratifyStep.value === 1) {
-			strataModel.value = null;
-		}
+	if (stratifyStep.value > 1) {
 		stratifyStep.value--;
+	} else {
+		strataModel.value = null;
 	}
 }
 
@@ -380,6 +399,8 @@ section {
 .buttons {
 	display: flex;
 	justify-content: space-between;
+	margin-left: auto;
+	gap: 1rem;
 }
 
 :deep(.p-button.p-button-outlined) {
@@ -391,10 +412,6 @@ section {
 	padding: 0 0.5rem 0 0.5rem;
 }
 
-.stratify-button {
-	margin-left: auto;
-}
-
 .output {
 	padding-left: 0.5rem;
 }
@@ -419,4 +436,12 @@ section {
 .add-model-config {
 	display: flex;
 }
+
+.nodes-require-types {
+	font-weight: var(--font-weight-semibold);
+}
+
+.buttons .p-button {
+	white-space: nowrap;
+}
 </style>
diff --git a/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts b/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
index c34dad840d..98c8784ac7 100644
--- a/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
+++ b/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
@@ -166,9 +166,6 @@ const DUMMY_VALUE = -999;
 export const convertToAMRModel = (g: IGraph<NodeData, EdgeData>) => g.amr;
 
 export const addState = (amr: Model, id: string, name: string) => {
-	if (amr.model.states.find((s) => s.id === id)) {
-		return;
-	}
 	amr.model.states.push({
 		id,
 		name,
@@ -188,9 +185,6 @@ export const addState = (amr: Model, id: string, name: string) => {
 };
 
 export const addTransition = (amr: Model, id: string, name: string, value?: number) => {
-	if (amr.model.transitions.find((t) => t.id === id)) {
-		return;
-	}
 	amr.model.transitions.push({
 		id,
 		input: [],
diff --git a/packages/client/hmi-client/src/services/models/stratification-service.ts b/packages/client/hmi-client/src/services/models/stratification-service.ts
index 0032262503..6af7a7daba 100644
--- a/packages/client/hmi-client/src/services/models/stratification-service.ts
+++ b/packages/client/hmi-client/src/services/models/stratification-service.ts
@@ -33,12 +33,10 @@ export async function fetchStratificationResult(modelA: string, modelB: string,
 }
 
 export function generateAgeStrataModel(stateNames: string[]): Model {
-	const states: State[] = stateNames.map((name, index) => ({
-		id: `A${index + 1}`,
+	const states: State[] = stateNames.map((name) => ({
+		id: name,
 		name,
-		description: `Number of individuals relative to the total population that are in age group A${
-			index + 1
-		}.`,
+		description: `Number of individuals relative to the total population that are in age group ${name}.`,
 		units: {
 			expression: 'person',
 			expression_mathml: '<ci>person</ci>'
@@ -161,12 +159,10 @@ export function generateAgeStrataModel(stateNames: string[]): Model {
 }
 
 export function generateLocationStrataModel(stateNames: string[]): Model {
-	const states: State[] = stateNames.map((name, index) => ({
-		id: `L${index + 1}`,
+	const states: State[] = stateNames.map((name) => ({
+		id: name,
 		name,
-		description: `Number of individuals relative to the total population that are in location L${
-			index + 1
-		}.`,
+		description: `Number of individuals relative to the total population that are in location ${name}.`,
 		units: {
 			expression: 'person',
 			expression_mathml: '<ci>person</ci>'
@@ -183,7 +179,7 @@ export function generateLocationStrataModel(stateNames: string[]): Model {
 					output: [innerState.id],
 					properties: {
 						name: `t&#832${i + 1}&#832${j + 1}`,
-						description: `Travel of an individual from location L${i + 1} and L${j + 1}.`
+						description: `Travel of an individual from location ${states[i].name} and ${states[j].name}.`
 					}
 				});
 			}
diff --git a/packages/client/hmi-client/tests/unit/services/stratification-service.spec.ts b/packages/client/hmi-client/tests/unit/services/stratification-service.spec.ts
index 5ddf8191ca..7c5b7d217e 100644
--- a/packages/client/hmi-client/tests/unit/services/stratification-service.spec.ts
+++ b/packages/client/hmi-client/tests/unit/services/stratification-service.spec.ts
@@ -21,20 +21,20 @@ describe('test generate age strata model', () => {
 			model: {
 				states: [
 					{
-						id: 'A1',
+						id: 'Young',
 						name: 'Young',
 						description:
-							'Number of individuals relative to the total population that are in age group A1.',
+							'Number of individuals relative to the total population that are in age group Young.',
 						units: {
 							expression: 'person',
 							expression_mathml: '<ci>person</ci>'
 						}
 					},
 					{
-						id: 'A2',
+						id: 'Old',
 						name: 'Old',
 						description:
-							'Number of individuals relative to the total population that are in age group A2.',
+							'Number of individuals relative to the total population that are in age group Old.',
 						units: {
 							expression: 'person',
 							expression_mathml: '<ci>person</ci>'
@@ -44,8 +44,8 @@ describe('test generate age strata model', () => {
 				transitions: [
 					{
 						id: 'c11',
-						input: ['A1', 'A1'],
-						output: ['A1', 'A1'],
+						input: ['Young', 'Young'],
+						output: ['Young', 'Young'],
 						properties: {
 							name: 'c&#8321&#8321',
 							description: 'Infective interaction between individuals.'
@@ -53,8 +53,8 @@ describe('test generate age strata model', () => {
 					},
 					{
 						id: 'c12',
-						input: ['A1', 'A2'],
-						output: ['A1', 'A2'],
+						input: ['Young', 'Old'],
+						output: ['Young', 'Old'],
 						properties: {
 							name: 'c&#8321&#8322',
 							description: 'Infective interaction between individuals.'
@@ -62,8 +62,8 @@ describe('test generate age strata model', () => {
 					},
 					{
 						id: 'c21',
-						input: ['A2', 'A1'],
-						output: ['A2', 'A1'],
+						input: ['Old', 'Young'],
+						output: ['Old', 'Young'],
 						properties: {
 							name: 'c&#8322&#8321',
 							description: 'Infective interaction between individuals.'
@@ -71,8 +71,8 @@ describe('test generate age strata model', () => {
 					},
 					{
 						id: 'c22',
-						input: ['A2', 'A2'],
-						output: ['A2', 'A2'],
+						input: ['Old', 'Old'],
+						output: ['Old', 'Old'],
 						properties: {
 							name: 'c&#8322&#8322',
 							description: 'Infective interaction between individuals.'
@@ -85,35 +85,36 @@ describe('test generate age strata model', () => {
 					rates: [
 						{
 							target: 'c11',
-							expression: 'A1*A1*beta11',
-							expression_mathml: '<apply><times/><ci>A1</ci><ci>A1</ci><ci>beta11</ci></apply>'
+							expression: 'Young*Young*beta11',
+							expression_mathml:
+								'<apply><times/><ci>Young</ci><ci>Young</ci><ci>beta11</ci></apply>'
 						},
 						{
 							target: 'c12',
-							expression: 'A1*A2*beta12',
-							expression_mathml: '<apply><times/><ci>A1</ci><ci>A2</ci><ci>beta12</ci></apply>'
+							expression: 'Young*Old*beta12',
+							expression_mathml: '<apply><times/><ci>Young</ci><ci>Old</ci><ci>beta12</ci></apply>'
 						},
 						{
 							target: 'c21',
-							expression: 'A2*A1*beta21',
-							expression_mathml: '<apply><times/><ci>A2</ci><ci>A1</ci><ci>beta21</ci></apply>'
+							expression: 'Old*Young*beta21',
+							expression_mathml: '<apply><times/><ci>Old</ci><ci>Young</ci><ci>beta21</ci></apply>'
 						},
 						{
 							target: 'c22',
-							expression: 'A2*A2*beta22',
-							expression_mathml: '<apply><times/><ci>A2</ci><ci>A2</ci><ci>beta22</ci></apply>'
+							expression: 'Old*Old*beta22',
+							expression_mathml: '<apply><times/><ci>Old</ci><ci>Old</ci><ci>beta22</ci></apply>'
 						}
 					],
 					initials: [
 						{
-							target: 'A1',
-							expression: 'A1init',
-							expression_mathml: '<ci>A1init</ci>'
+							target: 'Young',
+							expression: 'Younginit',
+							expression_mathml: '<ci>Younginit</ci>'
 						},
 						{
-							target: 'A2',
-							expression: 'A2init',
-							expression_mathml: '<ci>A2init</ci>'
+							target: 'Old',
+							expression: 'Oldinit',
+							expression_mathml: '<ci>Oldinit</ci>'
 						}
 					],
 					parameters: [
@@ -138,12 +139,12 @@ describe('test generate age strata model', () => {
 							value: 0.25
 						},
 						{
-							id: 'A1init',
+							id: 'Younginit',
 							description: "Proportion of population in age group 'Young' at timestep 0",
 							value: 0.5
 						},
 						{
-							id: 'A2init',
+							id: 'Oldinit',
 							description: "Proportion of population in age group 'Old' at timestep 0",
 							value: 0.5
 						}
@@ -179,8 +180,8 @@ describe('test generate age strata model', () => {
 						}
 					},
 					map: [
-						['A1', 'Pop'],
-						['A2', 'Pop'],
+						['Young', 'Pop'],
+						['Old', 'Pop'],
 						['c11', 'Infect'],
 						['c12', 'Infect'],
 						['c21', 'Infect'],
@@ -214,20 +215,20 @@ describe('test generate location strata model', () => {
 			model: {
 				states: [
 					{
-						id: 'L1',
+						id: 'Toronto',
 						name: 'Toronto',
 						description:
-							'Number of individuals relative to the total population that are in location L1.',
+							'Number of individuals relative to the total population that are in location Toronto.',
 						units: {
 							expression: 'person',
 							expression_mathml: '<ci>person</ci>'
 						}
 					},
 					{
-						id: 'L2',
+						id: 'Montreal',
 						name: 'Montreal',
 						description:
-							'Number of individuals relative to the total population that are in location L2.',
+							'Number of individuals relative to the total population that are in location Montreal.',
 						units: {
 							expression: 'person',
 							expression_mathml: '<ci>person</ci>'
@@ -237,20 +238,20 @@ describe('test generate location strata model', () => {
 				transitions: [
 					{
 						id: 't12',
-						input: ['L1'],
-						output: ['L2'],
+						input: ['Toronto'],
+						output: ['Montreal'],
 						properties: {
 							name: 't&#8321&#8322',
-							description: 'Travel of an individual from location L1 and L2.'
+							description: 'Travel of an individual from location Toronto and Montreal.'
 						}
 					},
 					{
 						id: 't21',
-						input: ['L2'],
-						output: ['L1'],
+						input: ['Montreal'],
+						output: ['Toronto'],
 						properties: {
 							name: 't&#8322&#8321',
-							description: 'Travel of an individual from location L2 and L1.'
+							description: 'Travel of an individual from location Montreal and Toronto.'
 						}
 					}
 				]
@@ -260,25 +261,25 @@ describe('test generate location strata model', () => {
 					rates: [
 						{
 							target: 't12',
-							expression: 'L1*tau12',
-							expression_mathml: '<apply><times/><ci>L1</ci><ci>tau12</ci></apply>'
+							expression: 'Toronto*tau12',
+							expression_mathml: '<apply><times/><ci>Toronto</ci><ci>tau12</ci></apply>'
 						},
 						{
 							target: 't21',
-							expression: 'L2*tau21',
-							expression_mathml: '<apply><times/><ci>L2</ci><ci>tau21</ci></apply>'
+							expression: 'Montreal*tau21',
+							expression_mathml: '<apply><times/><ci>Montreal</ci><ci>tau21</ci></apply>'
 						}
 					],
 					initials: [
 						{
-							target: 'L1',
-							expression: 'L1init',
-							expression_mathml: '<ci>L1init</ci>'
+							target: 'Toronto',
+							expression: 'Torontoinit',
+							expression_mathml: '<ci>Torontoinit</ci>'
 						},
 						{
-							target: 'L2',
-							expression: 'L2init',
-							expression_mathml: '<ci>L2init</ci>'
+							target: 'Montreal',
+							expression: 'Montrealinit',
+							expression_mathml: '<ci>Montrealinit</ci>'
 						}
 					],
 					parameters: [
@@ -293,12 +294,12 @@ describe('test generate location strata model', () => {
 							value: 0.5
 						},
 						{
-							id: 'L1init',
+							id: 'Torontoinit',
 							description: "Proportion of population in location 'Toronto' at timestep 0",
 							value: 0.5
 						},
 						{
-							id: 'L2init',
+							id: 'Montrealinit',
 							description: "Proportion of population in location 'Montreal' at timestep 0",
 							value: 0.5
 						}
@@ -334,8 +335,8 @@ describe('test generate location strata model', () => {
 						}
 					},
 					map: [
-						['L1', 'Pop'],
-						['L2', 'Pop'],
+						['Toronto', 'Pop'],
+						['Montreal', 'Pop'],
 						['t12', 'Strata'],
 						['t21', 'Strata']
 					]