From 5f7fad4f11e944f62a85c228e4a7378cee7cac55 Mon Sep 17 00:00:00 2001
From: dmccoystephenson <dmccoystephenson@gmail.com>
Date: Mon, 5 Sep 2022 16:52:32 +0000
Subject: [PATCH] Added some entity retrieval methods to the Environment class.

---
 src/environment.cpp      | 20 ++++++++++++++++++++
 src/grid.cpp             |  1 +
 src/header/entity.h      |  1 +
 src/header/environment.h |  4 +++-
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/environment.cpp b/src/environment.cpp
index 099ffcd..d334d20 100644
--- a/src/environment.cpp
+++ b/src/environment.cpp
@@ -52,4 +52,24 @@ namespace envlibcpp {
         std::cout << "=== " << getName() << " ===" << std::endl;
         std::cout << "Num entities: " << getNumEntities() << std::endl;
     }
+
+    envlibcpp::Entity& Environment::getFirstEntity() {
+        for (envlibcpp::Location &location : getGrid().getLocations()) {
+            if (location.getNumEntities() > 0) {
+                return location.getEntities()[0];
+            }
+        }
+        throw new std::exception();
+    }
+
+    envlibcpp::Entity& Environment::getEntity(int entityId) {
+        for (envlibcpp::Location &location : getGrid().getLocations()) {
+            for (envlibcpp::Entity &entity : location.getEntities()) {
+                if (entity.getId() == entityId) {
+                    return entity;
+                }
+            }
+        }
+        throw new std::exception();
+    }
 }
\ No newline at end of file
diff --git a/src/grid.cpp b/src/grid.cpp
index 4787c8e..e68fb0f 100644
--- a/src/grid.cpp
+++ b/src/grid.cpp
@@ -84,6 +84,7 @@ namespace envlibcpp {
         for (Location &location : locations) {
             if (location.isEntityPresent(entity)) {
                 location.removeEntity(entity);
+                entity.setLocationId(-1);
             }
         }
         entity.setGridId(-1);
diff --git a/src/header/entity.h b/src/header/entity.h
index 7ba070f..9fcf280 100644
--- a/src/header/entity.h
+++ b/src/header/entity.h
@@ -19,6 +19,7 @@ namespace envlibcpp {
             void setEnvironmentId(int id);
             void setGridId(int id);
             void setLocationId(int id);
+            envlibcpp::Entity& getEntity(int entityId);
         private:
             int id;
             std::string name;
diff --git a/src/header/environment.h b/src/header/environment.h
index 5bc561a..fee1be4 100644
--- a/src/header/environment.h
+++ b/src/header/environment.h
@@ -19,9 +19,11 @@ namespace envlibcpp {
             void addEntity(Entity &entity);
             void addEntityToLocation(Entity &entity, Location &location);
             void removeEntity(Entity &entity);
-            bool isEntityPresent(Entity &rntity);
+            bool isEntityPresent(Entity &entity);
             int getNumEntities();
             void printInfo();
+            envlibcpp::Entity& getFirstEntity();
+            envlibcpp::Entity& getEntity(int entityId);
         private:
             int id;
             std::string name;