diff --git a/geoplot.py b/geoplot.py
index 6f264d1..0cf8f61 100644
--- a/geoplot.py
+++ b/geoplot.py
@@ -1,293 +1,144 @@
-"""
-geoplot.py
-----------
-
-This visualization renders a 3-D plot of the data given the state
-trajectory of a simulation, and the path of the property to render.
-
-It generates an HTML file that contains code to render the plot
-using Cesium Ion, and the GeoJSON file of data provided to the plot.
-
-An example of its usage is as follows:
-
-```py
-from agent_torch.visualize import GeoPlot
-
-# create a simulation
-# ...
-
-# create a visualizer
-engine = GeoPlot(config, {
-  cesium_token: "...",
-  step_time: 3600,
-  coordinates = "agents/consumers/coordinates",
-  feature = "agents/consumers/money_spent",
-})
-
-# visualize in the runner-loop
-for i in range(0, num_episodes):
-  runner.step(num_steps_per_episode)
-  engine.render(runner.state_trajectory)
-```
-"""
-
 import re
 import json
-
 import pandas as pd
-import numpy as np
-
-from string import Template
-from agent_torch.core.helpers import get_by_path
+@@ -26,7 +25,6 @@
+# HTML template for Cesium visualization
+# This template defines the structure and behavior of the Cesium-based visualization.
+# It includes functions for interpolating colors, determining pixel sizes, and processing time-series data.
 
 geoplot_template = """
 <!doctype html>
 <html lang="en">
-	<head>
-		<meta charset="UTF-8" />
-		<meta
-			name="viewport"
-			content="width=device-width, initial-scale=1.0"
-		/>
-		<title>Cesium Time-Series Heatmap Visualization</title>
-		<script src="https://cesium.com/downloads/cesiumjs/releases/1.95/Build/Cesium/Cesium.js"></script>
-		<link
-			href="https://cesium.com/downloads/cesiumjs/releases/1.95/Build/Cesium/Widgets/widgets.css"
-			rel="stylesheet"
-		/>
-		<style>
-			#cesiumContainer {
-				width: 100%;
-				height: 100%;
-			}
-		</style>
-	</head>
-	<body>
-		<div id="cesiumContainer"></div>
-		<script>
-			// Your Cesium ion access token here
-			Cesium.Ion.defaultAccessToken = '$accessToken'
-
-			// Create the viewer
-			const viewer = new Cesium.Viewer('cesiumContainer')
-
-			function interpolateColor(color1, color2, factor) {
-				const result = new Cesium.Color()
-				result.red = color1.red + factor * (color2.red - color1.red)
-				result.green =
-					color1.green + factor * (color2.green - color1.green)
-				result.blue = color1.blue + factor * (color2.blue - color1.blue)
-				result.alpha = '$visualType' == 'size' ? 0.2 :
-					color1.alpha + factor * (color2.alpha - color1.alpha)
-				return result
-			}
-
-			function getColor(value, min, max) {
-				const factor = (value - min) / (max - min)
-				return interpolateColor(
-					Cesium.Color.BLUE,
-					Cesium.Color.RED,
-					factor
-				)
-			}
-
-			function getPixelSize(value, min, max) {
-				const factor = (value - min) / (max - min)
-				return 100 * (1 + factor)
-			}
-
-			function processTimeSeriesData(geoJsonData) {
-				const timeSeriesMap = new Map()
-				let minValue = Infinity
-				let maxValue = -Infinity
-
-				geoJsonData.features.forEach((feature) => {
-					const id = feature.properties.id
-					const time = Cesium.JulianDate.fromIso8601(
-						feature.properties.time
-					)
-					const value = feature.properties.value
-					const coordinates = feature.geometry.coordinates
-
-					if (!timeSeriesMap.has(id)) {
-						timeSeriesMap.set(id, [])
-					}
-					timeSeriesMap.get(id).push({ time, value, coordinates })
-
-					minValue = Math.min(minValue, value)
-					maxValue = Math.max(maxValue, value)
-				})
-
-				return { timeSeriesMap, minValue, maxValue }
-			}
-
-			function createTimeSeriesEntities(
-				timeSeriesData,
-				startTime,
-				stopTime
-			) {
-				const dataSource = new Cesium.CustomDataSource(
-					'AgentTorch Simulation'
-				)
-
-				for (const [id, timeSeries] of timeSeriesData.timeSeriesMap) {
-					const entity = new Cesium.Entity({
-						id: id,
-						availability: new Cesium.TimeIntervalCollection([
-							new Cesium.TimeInterval({
-								start: startTime,
-								stop: stopTime,
-							}),
-						]),
-						position: new Cesium.SampledPositionProperty(),
-						point: {
-							pixelSize: '$visualType' == 'size' ? new Cesium.SampledProperty(Number) : 10,
-							color: new Cesium.SampledProperty(Cesium.Color),
-						},
-						properties: {
-							value: new Cesium.SampledProperty(Number),
-						},
-					})
-
-					timeSeries.forEach(({ time, value, coordinates }) => {
-						const position = Cesium.Cartesian3.fromDegrees(
-							coordinates[0],
-							coordinates[1]
-						)
-						entity.position.addSample(time, position)
-						entity.properties.value.addSample(time, value)
-						entity.point.color.addSample(
-							time,
-							getColor(
-								value,
-								timeSeriesData.minValue,
-								timeSeriesData.maxValue
-							)
-						)
-
-						if ('$visualType' == 'size') {
-						  entity.point.pixelSize.addSample(
-  							time,
-  							getPixelSize(
-  								value,
-  								timeSeriesData.minValue,
-  								timeSeriesData.maxValue
-  							)
-  						)
-						}
-					})
-
-					dataSource.entities.add(entity)
-				}
-
-				return dataSource
-			}
-
-			// Example time-series GeoJSON data
-			const geoJsons = $data
-
-			const start = Cesium.JulianDate.fromIso8601('$startTime')
-			const stop = Cesium.JulianDate.fromIso8601('$stopTime')
-
-			viewer.clock.startTime = start.clone()
-			viewer.clock.stopTime = stop.clone()
-			viewer.clock.currentTime = start.clone()
-			viewer.clock.clockRange = Cesium.ClockRange.LOOP_STOP
-			viewer.clock.multiplier = 3600 // 1 hour per second
-
-			viewer.timeline.zoomTo(start, stop)
-
-			for (const geoJsonData of geoJsons) {
-				const timeSeriesData = processTimeSeriesData(geoJsonData)
-				const dataSource = createTimeSeriesEntities(
-					timeSeriesData,
-					start,
-					stop
-				)
-				viewer.dataSources.add(dataSource)
-				viewer.zoomTo(dataSource)
-			}
-		</script>
-	</body>
+@@ -41,9 +39,11 @@
+  <body>
+    <div id="cesiumContainer"></div>
+    <script>
+      // Setup Cesium viewer
+      Cesium.Ion.defaultAccessToken = '$accessToken';
+      const viewer = new Cesium.Viewer('cesiumContainer');
+      // Interpolate between two colors based on a factor
+      function interpolateColor(color1, color2, factor) {
+        const result = new Cesium.Color();
+        result.red = color1.red + factor * (color2.red - color1.red);
+@@ -53,16 +53,19 @@
+        return result;
+      }
+      // Get color based on value (between blue and red)
+      function getColor(value, min, max) {
+        const factor = (value - min) / (max - min);
+        return interpolateColor(Cesium.Color.BLUE, Cesium.Color.RED, factor);
+      }
+      // Get pixel size based on value
+      function getPixelSize(value, min, max) {
+        const factor = (value - min) / (max - min);
+        return 100 * (1 + factor);
+      }
+      // Process GeoJSON data to build a time-series map
+      function processTimeSeriesData(geoJsonData) {
+        const timeSeriesMap = new Map();
+        let minValue = Infinity;
+@@ -86,6 +89,7 @@
+        return { timeSeriesMap, minValue, maxValue };
+      }
+      // Create Cesium entities for the time series
+      function createTimeSeriesEntities(timeSeriesData, startTime, stopTime) {
+        const dataSource = new Cesium.CustomDataSource('AgentTorch Simulation');
+@@ -105,6 +109,7 @@
+            },
+          });
+          // Add each time sample
+          timeSeries.forEach(({ time, value, coordinates }) => {
+            const position = Cesium.Cartesian3.fromDegrees(coordinates[0], coordinates[1]);
+            entity.position.addSample(time, position);
+@@ -121,6 +126,7 @@
+        return dataSource;
+      }
+      // Load and visualize the time series data
+      const geoJsons = $data;
+      const start = Cesium.JulianDate.fromIso8601('$startTime');
+      const stop = Cesium.JulianDate.fromIso8601('$stopTime');
+@@ -133,6 +139,7 @@
+      viewer.timeline.zoomTo(start, stop);
+      // Load all GeoJSON datasets
+      for (const geoJsonData of geoJsons) {
+        const timeSeriesData = processTimeSeriesData(geoJsonData);
+        const dataSource = createTimeSeriesEntities(timeSeriesData, start, stop);
+@@ -144,17 +151,12 @@
 </html>
 """
 
+# Helper function to extract nested property from state based on path
+# This function uses the get_by_path utility to navigate nested dictionaries.
 
+# Helper function to extract nested property from the simulation state
 def read_var(state, var):
+    """Helper to extract nested property from state based on path."""
     return get_by_path(state, re.split("/", var))
 
+# GeoPlot class
+# This class encapsulates the logic for generating GeoJSON and HTML visualizations.
+# It takes configuration and visualization options as input.
 
+# GeoPlot class for generating visualization outputs
 class GeoPlot:
     def __init__(self, config, options):
-        self.config = config
-        (
-            self.cesium_token,
-            self.step_time,
-            self.entity_position,
-            self.entity_property,
-            self.visualization_type,
-        ) = (
-            options["cesium_token"],
-            options["step_time"],
-            options["coordinates"],
-            options["feature"],
+        """Initialize GeoPlot with config and visualization options."""
+@@ -173,24 +175,21 @@ def __init__(self, config, options):
             options["visualization_type"],
         )
 
+    # Render the trajectory to a GeoJSON and HTML visualization
+    # This method processes the state trajectory to generate GeoJSON features and an HTML file.
+    # It extracts coordinates and property values, generates timestamps, and constructs GeoJSON features.
+
     def render(self, state_trajectory):
+        """Render the trajectory to a GeoJSON and HTML visualization."""
+        """Render the trajectory to GeoJSON and HTML visualization."""
         coords, values = [], []
         name = self.config["simulation_metadata"]["name"]
         geodata_path, geoplot_path = f"{name}.geojson", f"{name}.html"
 
+        # Extract coordinates and property values from final states
+        # Extract final state coordinates and properties
         for i in range(0, len(state_trajectory) - 1):
             final_state = state_trajectory[i][-1]
-
             coords = np.array(read_var(final_state, self.entity_position)).tolist()
             values.append(
                 np.array(read_var(final_state, self.entity_property)).flatten().tolist()
             )
 
+        # Start time for the simulation
         start_time = pd.Timestamp.utcnow()
-        timestamps = [
-            start_time + pd.Timedelta(seconds=i * self.step_time)
-            for i in range(
-                self.config["simulation_metadata"]["num_episodes"]
-                * self.config["simulation_metadata"]["num_steps_per_episode"]
-            )
-        ]
+
+        # Generate timestamps spaced by step_time
+@@ -204,7 +203,7 @@ def render(self, state_trajectory):
 
         geojsons = []
+
+        # Construct GeoJSON features for each coordinate
+        # Construct GeoJSON features for visualization
         for i, coord in enumerate(coords):
             features = []
             for time, value_list in zip(timestamps, values):
-                features.append(
-                    {
-                        "type": "Feature",
-                        "geometry": {
-                            "type": "Point",
-                            "coordinates": [coord[1], coord[0]],
-                        },
-                        "properties": {
-                            "value": value_list[i],
-                            "time": time.isoformat(),
-                        },
-                    }
-                )
+@@ -221,14 +220,11 @@ def render(self, state_trajectory):
+                })
             geojsons.append({"type": "FeatureCollection", "features": features})
 
+        # Write GeoJSON file
+        # Write GeoJSON output file
         with open(geodata_path, "w", encoding="utf-8") as f:
             json.dump(geojsons, f, ensure_ascii=False, indent=2)
 
+        # Fill the HTML template with real data and token
+        # The HTML file is generated by substituting values into the Cesium template.
+        # It includes the Cesium token, start and stop times, GeoJSON data, and visualization type.
+
+        # Generate HTML file by substituting template values
         tmpl = Template(geoplot_template)
         with open(geoplot_path, "w", encoding="utf-8") as f:
             f.write(
-                tmpl.substitute(
-                    {
-                        "accessToken": self.cesium_token,
-                        "startTime": timestamps[0].isoformat(),
-                        "stopTime": timestamps[-1].isoformat(),
-                        "data": json.dumps(geojsons),
-                        "visualType": self.visualization_type,
-                    }
-                )
+@@ -240,3 +236,4 @@ def render(self, state_trajectory):
+                    "visualType": self.visualization_type,
+                })
             )
+
