diff --git a/src/AudibleInstruments.cpp b/src/AudibleInstruments.cpp
index b6f9ccb..1eabd8d 100644
--- a/src/AudibleInstruments.cpp
+++ b/src/AudibleInstruments.cpp
@@ -23,5 +23,5 @@ void init(rack::Plugin *p) {
 	p->addModel(modelVeils);
 	p->addModel(modelFrames);
 	// p->addModel(modelPeaks);
-	// p->addModel(modelStages);
+	p->addModel(modelStages);
 }
diff --git a/src/Stages.cpp b/src/Stages.cpp
index b7ca955..3c8f2b5 100644
--- a/src/Stages.cpp
+++ b/src/Stages.cpp
@@ -9,6 +9,58 @@
 static const int NUM_CHANNELS = 6;
 static const int BLOCK_SIZE = 8;
 
+struct GroupBuilder {
+
+	GroupBuilder() {
+		for (size_t i = 0; i < NUM_CHANNELS; i += 1) {
+			groupSize[i] = 0;
+		}
+	}
+
+	size_t groupSize[NUM_CHANNELS];
+	bool isPatched = false;
+
+	bool buildGroups(std::vector<Input> *inputs, size_t first, size_t count) {
+		bool changed = false;
+		isPatched = false;
+		size_t activeGroup = 0;
+
+		for (int i = count-1; i >= 0; i -= 1) {
+			auto patched = (*inputs)[first + i].active;
+
+			activeGroup += 1;
+			if (!patched) {
+				changed = changed || groupSize[i] != 0;
+				groupSize[i] = 0;
+			} else if (patched) {
+				isPatched = true;
+				changed = changed || groupSize[i] != activeGroup;
+				groupSize[i] = activeGroup;
+				activeGroup = 0;
+			}
+		}
+
+		return changed;
+	}
+
+	int groupForSegment(int segment) {
+		int group = 0;
+		int currentGroupSize = 0;
+
+		for (int i = 0; i < NUM_CHANNELS; i += 1) {
+			if (currentGroupSize <= 0) {
+				currentGroupSize = max(1, groupSize[i]);
+				group = i;
+			}
+
+			if (segment == i) {
+				return group;
+			}
+
+			currentGroupSize -= 1;
+		}
+	}
+};
 
 struct Stages : Module {
 	enum ParamIds {
@@ -33,8 +85,10 @@ struct Stages : Module {
 	};
 
 	stages::segment::Configuration configurations[NUM_CHANNELS];
+	bool configuration_changed[NUM_CHANNELS];
 	stages::SegmentGenerator segment_generator[NUM_CHANNELS];
 	stages::Oscillator oscillator[NUM_CHANNELS];
+	bool abloop;
 	// stages::ChainState chain_state;
 	// stages::Settings settings;
 
@@ -44,9 +98,11 @@ struct Stages : Module {
 
 	// Buffers
 	float envelopeBuffer[NUM_CHANNELS][BLOCK_SIZE] = {};
+	float typeLightsBuffer[NUM_CHANNELS][BLOCK_SIZE] = {};
 	stmlib::GateFlags last_gate_flags[NUM_CHANNELS] = {};
 	stmlib::GateFlags gate_flags[NUM_CHANNELS][BLOCK_SIZE] = {};
 	int blockIndex = 0;
+	GroupBuilder groupBuilder;
 
 	Stages() : Module(NUM_PARAMS, NUM_INPUTS, NUM_OUTPUTS, NUM_LIGHTS) {
 		// chain_state.Init(NULL, NULL);
@@ -54,12 +110,14 @@ struct Stages : Module {
 	}
 
 	void onReset() override {
+		abloop = false;
 		for (size_t i = 0; i < NUM_CHANNELS; ++i) {
 			segment_generator[i].Init();
 			oscillator[i].Init();
 
 			configurations[i].type = stages::segment::TYPE_RAMP;
 			configurations[i].loop = false;
+			configuration_changed[i] = true;
 		}
 	}
 
@@ -103,21 +161,52 @@ struct Stages : Module {
 			secondaries[i] = params[SHAPE_PARAMS + i].value;
 		}
 
-		// // Get patched state
-		// for (int i = 0; i < NUM_CHANNELS; i++) {
-		// 	bool p = inputs[GATE_INPUTS + i].active;
-		// 	if (p != patched[i])
-		// 		dirty_configurations[i] = true;
-		// 	patched[i] = p;
-		// }
+		// See if the group associations have changed since the last group
+		auto groups_changed = groupBuilder.buildGroups(&inputs, GATE_INPUTS, NUM_CHANNELS);
 
 		// Process block
 		stages::SegmentGenerator::Output out[BLOCK_SIZE] = {};
-		for (int i = 0; i < NUM_CHANNELS; i++) {
+		for (int i = 0; i < NUM_CHANNELS;) {
+			
+			// Check if the config needs applying to the segment generator for this group
+			bool segment_changed = groups_changed;
+			int numberOfLoops = 0;
+			for (size_t j = 0; j < max(1, groupBuilder.groupSize[i]); j += 1) {
+				numberOfLoops += configurations[i + j].loop ? 1 : 0;
+				segment_changed |= configuration_changed[i + j];
+				configuration_changed[i + j] = false;
+			}
+
+			if (segment_changed) {
+				if (numberOfLoops > 2) {
+					for (size_t j = 0; j < max(1, groupBuilder.groupSize[i]); j += 1) {
+						configurations[i + j].loop = false;
+					}
+				}
+				segment_generator[i].Configure(groupBuilder.groupSize[i] > 0, &configurations[i], max(1, groupBuilder.groupSize[i]));
+			}
+
+			// Set the segment parameters on the generator we're about to process
+			for (size_t j = 0; j < segment_generator[i].num_segments(); j += 1) {
+				segment_generator[i].set_segment_parameters(j, primaries[i + j], secondaries[i + j]);
+
+        oscillator[i + j].Render<stages::OSCILLATOR_SHAPE_TRIANGLE>(
+					0.00001f, 0.5f, typeLightsBuffer[i + j], BLOCK_SIZE
+				);
+			}
+
 			bool led_state = segment_generator[i].Process(gate_flags[i], out, BLOCK_SIZE);
+
+			// Set the outputs for the active segment in each output sample
+			// All outputs also go to the first segment
 			for (int j = 0; j < BLOCK_SIZE; j++) {
+				for (int k = 1; k < segment_generator[i].num_segments(); k += 1) {
+					envelopeBuffer[i + k][j] = k == out[j].segment ? out[j].value : 0.f;
+				}
 				envelopeBuffer[i][j] = out[j].value;
 			}
+
+			i += segment_generator[i].num_segments();
 		}
 	}
 
@@ -125,16 +214,50 @@ struct Stages : Module {
 		// Buttons
 		for (int i = 0; i < NUM_CHANNELS; i++) {
 			bool pressed = params[TYPE_PARAMS + i].value > 0.f;
-			if (pressed) {
+			if (pressed && pressedTime >= 0.f) {
 				pressedTime += engineGetSampleTime();
-			}
-			// Check if released
-			if (typeTriggers[i].process(!pressed)) {
 				if (pressedTime >= 1.f) {
+					pressedTime = -1.f;
+					configuration_changed[i] = true;
 					configurations[i].loop = !configurations[i].loop;
+
+					// ensure that we're the only loop item in the group
+					if (configurations[i].loop) {
+						int group = groupBuilder.groupForSegment(i);
+
+						if (abloop) {
+							// See how many loop items we have
+							int loopitems = 0;
+
+							for (int j = 0; j < groupBuilder.groupSize[group]; j += 1) {
+								loopitems += configurations[group + j].loop ? 1 : 0;
+							}
+
+							// Turn abloop off if we've got 2 or more loops
+							if (loopitems >= 2) {
+								abloop = false;
+							}
+
+							// If we've got >2 loops, clear down to the one loop
+							if (loopitems > 2) {
+								for (int j = 0; j < groupBuilder.groupSize[group]; j += 1) {
+									configurations[group + j].loop = (group + j) == i;
+								}
+							}
+						} else {
+							for (int j = 0; j < groupBuilder.groupSize[group]; j += 1) {
+								configurations[group + j].loop = (group + j) == i;
+							}
+						}
+					}
 				}
-				else {
+			}
+
+			// Check if released
+			if (typeTriggers[i].process(!pressed)) {
+				if (pressedTime >= 0.f) {
 					configurations[i].type = (stages::segment::Type) ((configurations[i].type + 1) % 3);
+					configuration_changed[i] = true;
 				}
 				pressedTime = 0.f;
 			}
@@ -154,13 +277,32 @@ struct Stages : Module {
 		}
 
 		// Output
+		int group = 0;
+		int currentGroupSize = 0;
+		int loopcount = 0;
 		for (int i = 0; i < NUM_CHANNELS; i++) {
 			float envelope = envelopeBuffer[i][blockIndex];
 			outputs[ENVELOPE_OUTPUTS + i].value = envelope * 10.f;
 			lights[ENVELOPE_LIGHTS + i].setBrightnessSmooth(envelope);
 
-			lights[TYPE_LIGHTS + i*2 + 0].setBrightness(configurations[i].type == 0 || configurations[i].type == 1);
-			lights[TYPE_LIGHTS + i*2 + 1].setBrightness(configurations[i].type == 1 || configurations[i].type == 2);
+			if (currentGroupSize <= 0) {
+				currentGroupSize = max(1, groupBuilder.groupSize[i]);
+				group = i;
+				loopcount = 0;
+			}
+			currentGroupSize -= 1;
+
+			loopcount += configurations[i].loop ? 1 : 0;
+			if (!configurations[i].loop) {
+				lights[TYPE_LIGHTS + i*2 + 0].setBrightness(configurations[i].type == 0 || configurations[i].type == 1);
+				lights[TYPE_LIGHTS + i*2 + 1].setBrightness(configurations[i].type == 1 || configurations[i].type == 2);
+			} else if (configurations[i].loop && loopcount == 1) {
+				lights[TYPE_LIGHTS + i*2 + 0].setBrightness((configurations[i].type == 0 || configurations[i].type == 1) * abs(typeLightsBuffer[i][blockIndex]));
+				lights[TYPE_LIGHTS + i*2 + 1].setBrightness((configurations[i].type == 1 || configurations[i].type == 2) * abs(typeLightsBuffer[i][blockIndex]));
+			} else {
+				lights[TYPE_LIGHTS + i*2 + 0].setBrightness((configurations[i].type == 0 || configurations[i].type == 1) * (1.f - abs(typeLightsBuffer[i][blockIndex])));
+				lights[TYPE_LIGHTS + i*2 + 1].setBrightness((configurations[i].type == 1 || configurations[i].type == 2) * (1.f - abs(typeLightsBuffer[i][blockIndex])));
+			}
 		}
 	}
 };
@@ -227,6 +369,22 @@ struct StagesWidget : ModuleWidget {
 		addChild(ModuleLightWidget::create<MediumLight<GreenLight>>(mm2px(Vec(48.07649, 103.19253)), module, Stages::ENVELOPE_LIGHTS + 4));
 		addChild(ModuleLightWidget::create<MediumLight<GreenLight>>(mm2px(Vec(59.51696, 103.19253)), module, Stages::ENVELOPE_LIGHTS + 5));
 	}
+
+	void appendContextMenu(Menu *menu) override {
+		Stages *module = dynamic_cast<Stages*>(this->module);
+
+		struct ABLoopItem : MenuItem {
+			Stages *module;
+			void onAction(EventAction &e) override {
+				module->abloop = true;
+			}
+		};
+
+		menu->addChild(MenuEntry::create());
+		ABLoopItem *abloopItem = MenuItem::create<ABLoopItem>("Set A/B Loop", CHECKMARK(module->abloop));
+		abloopItem->module = module;
+		menu->addChild(abloopItem);
+	}
 };