diff --git a/FlowEdge.cs b/FlowEdge.cs
index 5e5e251f8d7070394bab472d7c93250ac193a478..fd32617cb96e569508b7e2631f12df82f844aed8 100644
--- a/FlowEdge.cs
+++ b/FlowEdge.cs
@@ -5,17 +5,23 @@ using QuikGraph;
 public class FlowEdge : Edge<FlowNode>
 {
     public double MaxFlow { get; set; }
-    public double Residual { get; set; }
     
     public double CurrentFlow { get; set; }
+    
+    public double Residual
+    {
+        get => MaxFlow - CurrentFlow;
+        //set => CurrentFlow = MaxFlow - value;
+        //make the setter private
+        private set => CurrentFlow = MaxFlow - value;
+    }
     public bool IsBackwards { get; set; }
 
-    public FlowEdge(FlowNode source, FlowNode target, double maxFlow)
-        : base(source, target)
+    public FlowEdge(FlowNode source, FlowNode target, double maxFlow) : base(source, target)
     {
         IsBackwards = false;
         MaxFlow = maxFlow;
-        Residual = maxFlow;
+        CurrentFlow = 0;
     }
 
     public override string ToString()
diff --git a/FlowGraph.cs b/FlowGraph.cs
index 71e3195b797f0aec9c0f12aa74c3196915750546..8da4843a7e380db0c8c6f8d4c671883bf56294d1 100644
--- a/FlowGraph.cs
+++ b/FlowGraph.cs
@@ -36,6 +36,11 @@ public class FlowGraph
         var edge = new FlowEdge(sourceNode, targetNode, maxFlow);
         Graph.AddEdge(edge);
     }
+    
+    public void AddEdge(FlowEdge edge)
+    {
+        Graph.AddEdge(edge);
+    }
 
     public IEnumerable<FlowEdge> GetEdges()
     {
@@ -63,10 +68,11 @@ public class FlowGraph
             {
                 var clonedEdge = new FlowEdge(source, target, edge.MaxFlow)
                 {
-                    Residual = edge.Residual,
+                    //Residual = edge.Residual,
+                    CurrentFlow = edge.CurrentFlow,
                     IsBackwards = edge.IsBackwards
                 };
-                clonedGraph.Graph.AddEdge(clonedEdge);
+                clonedGraph.AddEdge(clonedEdge);
             }
         }
 
diff --git a/FordFulkersonAlgorithm.cs b/FordFulkersonAlgorithm.cs
index 3071aa8b6b3a9963f62752be4df3e609c9fe38a2..9908211b5aacfb9e63e91592792071777d160691 100644
--- a/FordFulkersonAlgorithm.cs
+++ b/FordFulkersonAlgorithm.cs
@@ -11,22 +11,26 @@ public class FordFulkersonAlgorithm
         BreadthFirstSearch,
         DepthFirstSearch
     }
-    
+
+    public readonly string SourceId;
+    public readonly string TargetId;
     private readonly FlowGraph _flowGraph;
     private readonly List<FlowGraph> _graphStates = new();
     
     public IReadOnlyList<FlowGraph> GraphStates => _graphStates;
 
-    public FordFulkersonAlgorithm(FlowGraph flowGraph)
+    public FordFulkersonAlgorithm(FlowGraph flowGraph, string sourceId, string targetId)
     {
         _flowGraph = flowGraph;
+        this.SourceId = sourceId;
+        this.TargetId = targetId;
     }
 
-    public double Run(string sourceId, string targetId, SearchStrategy strategy=SearchStrategy.BreadthFirstSearch)
+    public double Run(SearchStrategy strategy=SearchStrategy.BreadthFirstSearch)
     {
         // get source and target nodes
-        FlowNode? source = _flowGraph.GetVertexById(sourceId);
-        FlowNode? target = _flowGraph.GetVertexById(targetId);
+        FlowNode? source = _flowGraph.GetVertexById(SourceId);
+        FlowNode? target = _flowGraph.GetVertexById(TargetId);
 
         if (source == null || target == null)
             throw new ArgumentException("Invalid source or target node!");
@@ -34,14 +38,15 @@ public class FordFulkersonAlgorithm
         double maxFlow = 0.0;
         var pathFlow = new Dictionary<FlowEdge, double>();
 
-
         var allEdges = _flowGraph.GetEdges();
         foreach (var edge in allEdges)
         {
-            edge.Residual = edge.MaxFlow;
+            edge.CurrentFlow = 0;
         }
         
         Console.WriteLine("Start of Ford-Fulkerson Algorithm...");
+        
+        _graphStates.Clear();
         SaveGraphState();
 
         // execute as long as there is an augmenting path
@@ -60,7 +65,8 @@ public class FordFulkersonAlgorithm
             // add forward flow along the augmenting path
             foreach (var edge in pathFlow.Keys)
             {
-                edge.Residual -= pathMinFlow;
+                //edge.Residual -= pathMinFlow;
+                edge.CurrentFlow += pathMinFlow;
                 Console.WriteLine($"Updating forward edge {edge.Source.Id} -> {edge.Target.Id}, New Residual: {edge.Residual}");
             }
             
@@ -70,23 +76,21 @@ public class FordFulkersonAlgorithm
             foreach (var edge in pathFlow.Keys)
             {
                 // check if backwards edge exists
-                FlowEdge? reverseEdge;
-                if (!_flowGraph.Graph.TryGetEdge(edge.Target, edge.Source, out reverseEdge))
+                if (!_flowGraph.Graph.TryGetEdge(edge.Target, edge.Source, out var reverseEdge))
                 {
                     // create a backwards edge if it doesn't exist
-                    reverseEdge = new FlowEdge(edge.Target, edge.Source, edge.MaxFlow);
+                    reverseEdge = new FlowEdge(edge.Target, edge.Source, pathMinFlow);
 
-                    reverseEdge.Residual = 0;
+                    reverseEdge.CurrentFlow = 0;
                     reverseEdge.IsBackwards = true;
                     
                     _flowGraph.Graph.AddEdge(reverseEdge);
                     
                     Console.WriteLine($"Creating reverse edge {reverseEdge.Source.Id} -> {reverseEdge.Target.Id}");
                 }
-
-                if (reverseEdge != null)
+                else if (reverseEdge != null)
                 {
-                    reverseEdge.Residual += pathMinFlow;
+                    reverseEdge.MaxFlow += pathMinFlow;
                     
                     Console.WriteLine($"Updating reverse edge {reverseEdge.Source.Id} -> {reverseEdge.Target.Id}, New Residual: {reverseEdge.Residual}");
                 }
diff --git a/GraphVisualizer.cs b/GraphVisualizer.cs
index e823878879931c0bcc702c7f7f4f3761a5b6c9df..27de5483a02397ebe61f38ecb836bbba9c884f12 100644
--- a/GraphVisualizer.cs
+++ b/GraphVisualizer.cs
@@ -4,7 +4,7 @@ namespace FlowForge;
 
 public static class GraphVisualizer
 {
-    public static Graph ConvertToMsaglGraph(FlowGraph? flowGraph)
+    public static Graph ConvertToMsaglGraph(FlowGraph flowGraph)
     {
         Graph msaglGraph = new Graph();
 
@@ -22,12 +22,12 @@ public static class GraphVisualizer
 
             if (edge.IsBackwards)
             {
-                msaglEdge.LabelText = $"{0}/{edge.Residual}";
+                msaglEdge.LabelText = $"{edge.CurrentFlow}/{edge.MaxFlow}";
                 msaglEdge.Attr.Color = Color.Gray;
             }
             else
             {
-                msaglEdge.LabelText = $"{edge.MaxFlow-edge.Residual}/{edge.MaxFlow}";
+                msaglEdge.LabelText = $"{edge.CurrentFlow}/{edge.MaxFlow}";
 
                 if (edge.MaxFlow-edge.Residual > 0)
                 {
diff --git a/MainWindow.xaml.cs b/MainWindow.xaml.cs
index c0044b1e3f4e67101bf362a2630a9ef6407f73cf..4ebddf39caa0214d64a1a3f8caba03b52213eb0d 100644
--- a/MainWindow.xaml.cs
+++ b/MainWindow.xaml.cs
@@ -21,13 +21,13 @@ namespace FlowForge
     public partial class MainWindow : Window
     {
         private GViewer? _gViewer;
-        private FlowGraph? _flowGraph;
+        private FlowGraph _flowGraph = new FlowGraph();
         private FordFulkersonAlgorithm? _fordFulkerson;
 
-        private readonly Dictionary<string, Microsoft.Msagl.Core.Geometry.Point> _nodePositions =
-            new Dictionary<string, Microsoft.Msagl.Core.Geometry.Point>();
+        private readonly Dictionary<string, Microsoft.Msagl.Core.Geometry.Point> _nodePositions = new Dictionary<string, Microsoft.Msagl.Core.Geometry.Point>();
 
-        private MsaglDrawing.Graph _msaglGraph;
+        private MsaglDrawing.Graph _msaglGraph = new MsaglDrawing.Graph();
+        bool _isInitialized = false;
 
         public MainWindow()
         {
@@ -72,9 +72,9 @@ namespace FlowForge
 
             //SaveNodePositions(_gViewer.Graph); // Save positions before running the algorithm
 
-            double maxFlow = _fordFulkerson.Run("1", "4", FordFulkersonAlgorithm.SearchStrategy.DepthFirstSearch);
-
+            double maxFlow = _fordFulkerson.Run(FordFulkersonAlgorithm.SearchStrategy.DepthFirstSearch);
 
+            _isInitialized = true;
             DisplayGraph();
             MessageBox.Show($"Maximum flow from source to sink: {maxFlow}");
         }
@@ -90,8 +90,9 @@ namespace FlowForge
 
             if (_gViewer != null) _gViewer.NeedToCalculateLayout = true;
             SaveNodePositions(_gViewer?.Graph); // Save positions before running the algorithm
-            double maxFlow = _fordFulkerson.Run("1", "4", FordFulkersonAlgorithm.SearchStrategy.BreadthFirstSearch);
+            double maxFlow = _fordFulkerson.Run(FordFulkersonAlgorithm.SearchStrategy.BreadthFirstSearch);
 
+            _isInitialized = true;
             DisplayGraph();
             MessageBox.Show($"Maximum flow from source to sink: {maxFlow}");
         }
@@ -121,7 +122,7 @@ namespace FlowForge
             _flowGraph.AddEdge("1", "5", 200);
 
             // Initialize Ford-Fulkerson algorithm with the flow graph
-            _fordFulkerson = new FordFulkersonAlgorithm(_flowGraph);
+            _fordFulkerson = new FordFulkersonAlgorithm(_flowGraph, "1", "4");
         }
 
         private void DisplayGraph()
@@ -172,7 +173,7 @@ namespace FlowForge
                     sugiyamaSettings.EdgeRoutingSettings.Padding = 1;
                     sugiyamaSettings.RepetitionCoefficientForOrdering = 100;
                     sugiyamaSettings.EdgeRoutingSettings.KeepOriginalSpline = false;
-                    
+
                     Microsoft.Msagl.Miscellaneous.LayoutHelpers.RouteAndLabelEdges(_msaglGraph.GeometryGraph,
                         sugiyamaSettings, _msaglGraph.GeometryGraph.Edges, 100, new CancelToken());
 
@@ -237,7 +238,6 @@ namespace FlowForge
                 var sourceCenter = edge.Source.Center;
                 var targetCenter = edge.Target.Center;
 
-                // Create a straight line from source to target
                 // Create a straight line from source to target
                 edge.Curve = new LineSegment(sourceCenter, targetCenter);
                 edge.LineWidth = 2;
@@ -249,16 +249,17 @@ namespace FlowForge
                 edge.Attr.ArrowheadLength = 1;
             }
         }
-        
+
         private void ExportToLatexButton_Click(object sender, RoutedEventArgs e)
         {
-            if (_flowGraph == null)
+            if (!_isInitialized)
             {
                 MessageBox.Show("Please run an algorithm first.", "Error", MessageBoxButton.OK, MessageBoxImage.Error);
                 return;
             }
 
-            TikzCodeGenerator tikzCodeGenerator = new TikzCodeGenerator(_flowGraph, _nodePositions, _fordFulkerson.GraphStates);
+            if (_fordFulkerson == null) return;
+            TikzCodeGenerator tikzCodeGenerator = new TikzCodeGenerator(_flowGraph, _nodePositions, _fordFulkerson);
             tikzCodeGenerator.ExportGraphToLatex(IncludeAnimationCheckbox.IsChecked, StandaloneCheckbox.IsChecked);
         }
     }
diff --git a/TikzCodeGenerator.cs b/TikzCodeGenerator.cs
index 6021d3094209dd6062be4f525ccd8744f3ec8b11..33ac032896180e9635f2b80317e77be76751baa2 100644
--- a/TikzCodeGenerator.cs
+++ b/TikzCodeGenerator.cs
@@ -10,13 +10,15 @@ public class TikzCodeGenerator
 {
     private readonly FlowGraph _flowGraph;
     private readonly Dictionary<string, Microsoft.Msagl.Core.Geometry.Point> _nodePositions;
-    private IReadOnlyList<FlowGraph> _graphStates;
+    private readonly FordFulkersonAlgorithm _fordFulkerson;
+    private readonly IReadOnlyList<FlowGraph> _graphStates;
 
-    public TikzCodeGenerator(FlowGraph flowGraph, Dictionary<string, Microsoft.Msagl.Core.Geometry.Point> nodePositions, IReadOnlyList<FlowGraph> graphStates)
+    public TikzCodeGenerator(FlowGraph flowGraph, Dictionary<string, Microsoft.Msagl.Core.Geometry.Point> nodePositions, FordFulkersonAlgorithm fordFulkerson)
     {
         this._flowGraph = flowGraph;
         this._nodePositions = nodePositions;
-        this._graphStates = graphStates;
+        this._fordFulkerson = fordFulkerson;
+        this._graphStates = _fordFulkerson.GraphStates;
     }
 
     private string GenerateTikzCode_Curved(bool includeAnimations, bool isStandalone)
@@ -44,37 +46,80 @@ public class TikzCodeGenerator
         {
             BuildTikzEdges(tikzBuilder);
         }
-
-        if (isStandalone)
-        {
-            BuildTikzFooter(tikzBuilder);
-        }
+        
+        BuildTikzFooter(tikzBuilder, isStandalone);
 
         return tikzBuilder.ToString();
     }
 
-    private static void BuildTikzFooter(StringBuilder tikzBuilder)
+    private static void BuildTikzHeader(StringBuilder tikzBuilder)
     {
-        tikzBuilder.AppendLine("\\end{tikzpicture}");
-        tikzBuilder.AppendLine("\\end{frame} ");
-        tikzBuilder.AppendLine("\\end{document}");
+        tikzBuilder.AppendLine("\\documentclass[tikz,border=3mm]{beamer}");
+        
+        tikzBuilder.AppendLine("\\usepackage{tikz}");
+        tikzBuilder.AppendLine("\\usetikzlibrary{arrows.meta,positioning}");
+        
+        tikzBuilder.AppendLine("\\usepackage{amsmath}");
+        
+        tikzBuilder.AppendLine("\\begin{document}");
+        tikzBuilder.AppendLine("\\begin{frame}{Graph Animation}");
     }
+    
+    private static void BuildTikzSettings(StringBuilder tikzBuilder)
+    {
+        // Global scaling and styles for smaller nodes and text
+        tikzBuilder.AppendLine(
+            "\\begin{tikzpicture}[scale=2,->,>=stealth,shorten >=1pt,auto,node distance=2cm,thick,");
+        tikzBuilder.AppendLine("main node/.style={circle,draw,minimum size=5mm,inner sep=1pt,font=\\scriptsize}]");
+    }
+    
+    private void BuildTikzNodes(StringBuilder tikzBuilder)
+    {
+        // Nodes
+        foreach (var vertex in _flowGraph.Graph.Vertices)
+        {
+            if (_nodePositions.ContainsKey(vertex.Id))
+            {
+                var position = _nodePositions[vertex.Id];
+                // Format x and y as floating-point values with two decimal places
+                string formattedX =
+                    (position.X / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
+                string formattedY =
+                    (position.Y / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
 
+                // Determine node color
+                string colorStyle = "circle,draw";
+                if (vertex.Id == _fordFulkerson.SourceId)
+                {
+                    colorStyle = "circle,draw,fill=green!50";
+                }
+                else if (vertex.Id == _fordFulkerson.TargetId)
+                {
+                    colorStyle = "circle,draw,fill=red!50";
+                }
+                
+                // Write the TikZ node with the appropriate style
+                tikzBuilder.AppendLine($"\\node[main node, {colorStyle}] ({vertex.Id}) at ({formattedX},{formattedY}) {{{vertex.Id}}};");
+            }
+        }
+    }
+    
     private void BuildTikzEdges(StringBuilder tikzBuilder)
     {
         // Edges
         foreach (var edge in _flowGraph.GetEdges())
         {
+            string label = $"{edge.CurrentFlow}/{edge.MaxFlow}";
+            
             if (edge.IsBackwards)
             {
-                string label = $"0/{edge.Residual}";
+                
                 // Reverse edge styling (curved, dashed, different color)
                 tikzBuilder.AppendLine(
                     $"\\path[every node/.style={{font=\\tiny}},bend right,dashed,red] ({edge.Source.Id}) edge [pos=0.5,sloped] node {{\\texttt{{{label}}}}} ({edge.Target.Id});");
             }
             else
             {
-                string label = $"{edge.MaxFlow - edge.Residual}/{edge.MaxFlow}";
                 // Normal edge styling
                 tikzBuilder.AppendLine(
                     $"\\path[every node/.style={{font=\\tiny}}] ({edge.Source.Id}) edge [pos=0.5,sloped] node {{\\texttt{{{label}}}}} ({edge.Target.Id});");
@@ -96,7 +141,7 @@ public class TikzCodeGenerator
 
             foreach (var edge in stateGraph.GetEdges())
             {
-                string label = $"{edge.Residual}/{edge.MaxFlow}";
+                string label = $"{edge.CurrentFlow}/{edge.MaxFlow}";
 
                 if (edge.IsBackwards)
                 {
@@ -116,48 +161,18 @@ public class TikzCodeGenerator
             tikzBuilder.AppendLine("}");
         }
     }
-
-    private void BuildTikzNodes(StringBuilder tikzBuilder)
+    
+    private static void BuildTikzFooter(StringBuilder tikzBuilder, bool isStandalone)
     {
-        // Nodes
-        foreach (var vertex in _flowGraph.Graph.Vertices)
+        tikzBuilder.AppendLine("\\end{tikzpicture}");
+        
+        if (!isStandalone)
         {
-            if (_nodePositions.ContainsKey(vertex.Id))
-            {
-                var position = _nodePositions[vertex.Id];
-                // Format x and y as floating-point values with two decimal places
-                string formattedX =
-                    (position.X / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
-                string formattedY =
-                    (position.Y / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
-
-                tikzBuilder.AppendLine(
-                    $"\\node[main node] ({vertex.Id}) at ({formattedX},{formattedY}) {{{vertex.Id}}};");
-            }
+            tikzBuilder.AppendLine("\\end{frame} ");
+            tikzBuilder.AppendLine("\\end{document}");
         }
     }
-
-    private static void BuildTikzSettings(StringBuilder tikzBuilder)
-    {
-        // Global scaling and styles for smaller nodes and text
-        tikzBuilder.AppendLine(
-            "\\begin{tikzpicture}[scale=3,->,>=stealth,shorten >=1pt,auto,node distance=2cm,thick,");
-        tikzBuilder.AppendLine("main node/.style={circle,draw,minimum size=5mm,inner sep=1pt,font=\\scriptsize}]");
-    }
-
-    private static void BuildTikzHeader(StringBuilder tikzBuilder)
-    {
-        tikzBuilder.AppendLine("\\documentclass[tikz,border=3mm]{beamer}");
-        
-        tikzBuilder.AppendLine("\\usepackage{tikz}");
-        tikzBuilder.AppendLine("\\usetikzlibrary{arrows.meta,positioning}");
-        
-        tikzBuilder.AppendLine("\\usepackage{amsmath}");
-        
-        tikzBuilder.AppendLine("\\begin{document}");
-        tikzBuilder.AppendLine("\\begin{frame}{Graph Animation}");
-    }
-
+   
     public void ExportGraphToLatex(bool? includeAnimation, bool? standalone)
     {
         bool includeAnimationChecked = includeAnimation ?? false;
@@ -191,74 +206,4 @@ public class TikzCodeGenerator
             MessageBox.Show($"Error exporting graph: {ex.Message}", "Export Failed", MessageBoxButton.OK, MessageBoxImage.Error);
         }
     }
-
-    private string GenerateTikzCode_WithAnimation()
-    {
-        throw new NotImplementedException();
-    }
-
-    private string StripStandaloneWrapper(string tikzCode)
-    {
-        var lines = tikzCode.Split('\n').ToList();
-        lines.RemoveAll(line => line.StartsWith("\\documentclass") || line.StartsWith("\\usepackage") || line.StartsWith("\\begin{document}") || line.StartsWith("\\end{document}"));
-        return string.Join("\n", lines);
-    }
-    
-    private string GenerateTikzCode()
-    {
-        if (_flowGraph == null)
-            throw new InvalidOperationException("Flow graph is not initialized.");
-
-        var tikzBuilder = new System.Text.StringBuilder();
-        tikzBuilder.AppendLine("\\documentclass[tikz,border=3mm]{standalone}");
-        tikzBuilder.AppendLine("\\usetikzlibrary{arrows.meta,positioning}");
-        tikzBuilder.AppendLine("\\usepackage{amsmath}");
-        tikzBuilder.AppendLine("\\begin{document}");
-
-        // Global scaling and styles for smaller nodes and text
-        tikzBuilder.AppendLine(
-            "\\begin{tikzpicture}[scale=3,->,>=stealth,shorten >=1pt,auto,node distance=2cm,thick,");
-        tikzBuilder.AppendLine("main node/.style={circle,draw,minimum size=1mm,inner sep=1pt,font=\\scriptsize}]");
-
-        // Nodes
-        foreach (var vertex in _flowGraph.Graph.Vertices)
-        {
-            if (_nodePositions.ContainsKey(vertex.Id))
-            {
-                var position = _nodePositions[vertex.Id];
-                // Format x and y as floating-point values with two decimal places
-                string formattedX =
-                    (position.X / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
-                string formattedY =
-                    (position.Y / 100.0).ToString("F2", System.Globalization.CultureInfo.InvariantCulture);
-
-                tikzBuilder.AppendLine(
-                    $"\\node[main node] ({vertex.Id}) at ({formattedX},{formattedY}) {{{vertex.Id}}};");
-            }
-        }
-
-        // Edges
-        foreach (var edge in _flowGraph.GetEdges())
-        {
-            string label = $"{edge.Residual}/{edge.MaxFlow}";
-
-            if (edge.IsBackwards)
-            {
-                // Reverse edge styling: offset slightly above forward edges
-                tikzBuilder.AppendLine(
-                    $"\\path[every node/.style={{font=\\tiny}},dashed,red] ({edge.Source.Id}) edge [pos=0.5,above] node {{\\texttt{{{label}}}}} ({edge.Target.Id});");
-            }
-            else
-            {
-                // Normal edge styling
-                tikzBuilder.AppendLine(
-                    $"\\path[every node/.style={{font=\\tiny}}] ({edge.Source.Id}) edge [pos=0.5,below] node {{\\texttt{{{label}}}}} ({edge.Target.Id});");
-            }
-        }
-
-        tikzBuilder.AppendLine("\\end{tikzpicture}");
-        tikzBuilder.AppendLine("\\end{document}");
-
-        return tikzBuilder.ToString();
-    }
 }
\ No newline at end of file