3 files changed, 30 insertions, 7 deletions

diff --git a/subprojects/store-query-viatra/build.gradle.kts b/subprojects/store-query-viatra/build.gradle.kts
index 1f804247..ef73e10a 100644
--- a/subprojects/store-query-viatra/build.gradle.kts
+++ b/subprojects/store-query-viatra/build.gradle.kts
@@ -10,7 +10,6 @@ plugins {
 
 dependencies {
         implementation(libs.ecore)
-        implementation ("guru.nidi:graphviz-java:0.18.1")
         api(libs.viatra)
         api(project(":refinery-store-query"))
         api(project(":refinery-store-reasoning"))
diff --git a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/dse/internal/DesignSpaceExplorationAdapterImpl.java b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/dse/internal/DesignSpaceExplorationAdapterImpl.java
index c4c51e79..0fcc719d 100644
--- a/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/dse/internal/DesignSpaceExplorationAdapterImpl.java
+++ b/subprojects/store-query-viatra/src/main/java/tools/refinery/store/query/dse/internal/DesignSpaceExplorationAdapterImpl.java
@@ -15,6 +15,7 @@ import tools.refinery.store.query.resultset.ResultSet;
 import tools.refinery.store.representation.Symbol;
 import tools.refinery.store.tuple.Tuple;
 import tools.refinery.store.tuple.Tuple1;
+import tools.refinery.visualization.ModelVisualizerAdapter;
 
 import java.util.*;
 
@@ -40,6 +41,8 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         private Random random = new Random();
         private Objective[][] leveledObjectives;
         private boolean isNewState = false;
+        private final boolean isVisualizationEnabled;
+        private final ModelVisualizerAdapter modelVisualizerAdapter;
 
         public List<Long> getTrajectory() {
                 return new LinkedList<>(trajectory);
@@ -66,6 +69,8 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                 statesAndUntraversedActivations = new HashMap<>();
                 statesAndTraversedActivations = new HashMap<>();
                 strategy = storeAdapter.getStrategy();
+                modelVisualizerAdapter = model.tryGetAdapter(ModelVisualizerAdapter.class).orElse(null);
+                isVisualizationEnabled = modelVisualizerAdapter != null;
 
         }
 
@@ -129,6 +134,10 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                 if (trajectory.size() < 2) {
                         return false;
                 }
+                if (isVisualizationEnabled) {
+                        modelVisualizerAdapter.addTransition(trajectory.get(trajectory.size() - 1),
+                                        trajectory.get(trajectory.size() - 2), "backtrack");
+                }
                 model.restore(trajectory.get(trajectory.size() - 2));
                 trajectory.remove(trajectory.size() - 1);
                 return true;
@@ -137,6 +146,10 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         @Override
         public void restoreTrajectory(List<Long> trajectory) {
                 model.restore(trajectory.get(trajectory.size() - 1));
+//              if (isVisualizationEnabled) {
+//                      modelVisualizerAdapter.addTransition(this.trajectory.get(trajectory.size() - 1),
+//                                      trajectory.get(trajectory.size() - 1), "restore");
+//              }
                 this.trajectory = trajectory;
 
         }
@@ -173,6 +186,9 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         public void newSolution() {
                 var state = model.getState();
                 solutions.add(state);
+                if (isVisualizationEnabled) {
+                        modelVisualizerAdapter.addSolution(state);
+                }
         }
 
         @Override
@@ -212,6 +228,11 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                 isNewState = !statesAndUntraversedActivations.containsKey(newState);
                 statesAndUntraversedActivations.put(newState, getAllActivations());
                 statesAndTraversedActivations.put(newState, new HashSet<>());
+                if (isVisualizationEnabled) {
+                        modelVisualizerAdapter.addTransition(trajectory.get(trajectory.size() - 2),
+                                        trajectory.get(trajectory.size() - 1), activation.transformationRule().getName(),
+                                        activation.activation());
+                }
                 return true;
         }
 
diff --git a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/dse/DebugTest.java b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/dse/DebugTest.java
index 969afbcb..821be7e6 100644
--- a/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/dse/DebugTest.java
+++ b/subprojects/store-query-viatra/src/test/java/tools/refinery/store/query/dse/DebugTest.java
@@ -5,6 +5,7 @@ import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.query.ModelQueryAdapter;
 import tools.refinery.store.query.dnf.Query;
 import tools.refinery.store.query.dse.internal.TransformationRule;
+import tools.refinery.store.query.dse.strategy.BestFirstStrategy;
 import tools.refinery.store.query.dse.strategy.DepthFirstStrategy;
 import tools.refinery.store.query.viatra.ViatraModelQueryAdapter;
 import tools.refinery.store.query.view.AnySymbolView;
@@ -12,6 +13,7 @@ import tools.refinery.store.query.view.KeyOnlyView;
 import tools.refinery.store.representation.Symbol;
 import tools.refinery.store.tuple.Tuple;
 import tools.refinery.visualization.ModelVisualizerAdapter;
+import tools.refinery.visualization.internal.FileFormat;
 
 public class DebugTest {
         private static final Symbol<Boolean> classModel = Symbol.of("ClassModel", 1);
@@ -83,30 +85,31 @@ public class DebugTest {
                                 .symbols(classModel, classElement, feature, isEncapsulatedBy, encapsulates, classes, features)
                                 .with(ViatraModelQueryAdapter.builder()
                                                 .queries(createClassPrecondition, createFeaturePrecondition))
+                                .with(ModelVisualizerAdapter.builder())
                                 .with(DesignSpaceExplorationAdapter.builder()
                                                 .transformations(createClassRule, createFeatureRule)
                                                 .strategy(new DepthFirstStrategy(4).continueIfHardObjectivesFulfilled()
 //                                              .strategy(new BestFirstStrategy(4).continueIfHardObjectivesFulfilled()
 //                                                              .goOnOnlyIfFitnessIsBetter()
                                                 ))
-                                .with(ModelVisualizerAdapter.builder())
                                 .build();
 
                 var model = store.createEmptyModel();
                 var dseAdapter = model.getAdapter(DesignSpaceExplorationAdapter.class);
+//              dseAdapter.setRandom(1);
                 var queryEngine = model.getAdapter(ModelQueryAdapter.class);
 
                 var modelElementInterpretation = model.getInterpretation(classModel);
-                modelElementInterpretation.put(dseAdapter.createObject(), true);
+                var classElementInterpretation = model.getInterpretation(classElement);
+                var modelElement = dseAdapter.createObject();
+                modelElementInterpretation.put(modelElement, true);
+                classElementInterpretation.put(modelElement, true);
                 queryEngine.flushChanges();
 
 
                 var states = dseAdapter.explore();
                 var visualizer = model.getAdapter(ModelVisualizerAdapter.class);
-                for (var state : states) {
-                        var visualization = visualizer.createVisualizationForModelState(state);
-                        visualizer.saveVisualization(visualization, "test_output" + state.hashCode() + ".png");
-                }
+                visualizer.renderDesignSpace("test_output", FileFormat.SVG);
                 System.out.println("states size: " + states.size());
                 System.out.println("states: " + states);