diff --git a/lib/main.dart b/lib/main.dart
index fc6b18c..85f2cf1 100644
--- a/lib/main.dart
+++ b/lib/main.dart
@@ -212,12 +212,12 @@ class _MainPageState extends State<MainPage> with WidgetsBindingObserver {
     List<Widget> result = [];
     if (view.isThereBadges) {
       result.add(
-            Wrap(
-              alignment: WrapAlignment.center,
-              spacing: 10.0,
-              runSpacing: 4.0,
-              children: _buildBadges(view.badges),
-            )
+        Wrap(
+          alignment: WrapAlignment.center,
+          spacing: 12.0,
+          runSpacing: 4.0,
+          children: _buildBadges(view.badges),
+        )
       );
 
     }
diff --git a/lib/ui_builder_class.dart b/lib/ui_builder_class.dart
index f7e0648..651defc 100644
--- a/lib/ui_builder_class.dart
+++ b/lib/ui_builder_class.dart
@@ -3,7 +3,7 @@ part of 'main.dart';
 class UIBuilder {
   EntityCollection _entities;
   Map<String, View> _views;
-  List _topBadgeDomains = ["alarm_control_panel", "binary_sensor", "device_tracker", "updater", "sun", "timer", "sensor"];
+  static List badgeDomains = ["alarm_control_panel", "binary_sensor", "device_tracker", "updater", "sun", "timer", "sensor"];
 
   bool get isEmpty => _views.length == 0;
   Map<String, View> get views => _views ?? {};
@@ -12,6 +12,10 @@ class UIBuilder {
     _views = {};
   }
 
+  static bool isBadge(String domain) {
+    return badgeDomains.contains(domain);
+  }
+
   void build(EntityCollection entitiesCollection) {
     _entities = entitiesCollection;
     _views.clear();
@@ -36,18 +40,7 @@ class UIBuilder {
         viewGroupEntity.childEntities.forEach((
               entityId) { //Each entity or group in view
           if (_entities.isExist(entityId)) {
-            Entity entityToAdd = _entities.get(entityId);
-            if (!entityToAdd.isGroup) {
-              if (_topBadgeDomains.contains(entityToAdd.domain)) {
-                //This is badge
-                view.addBadge(entityId);
-              } else {
-                //This is a standalone entity
-                view.addEntityWithoutGroup(entityToAdd);
-              }
-            } else {
-              view.addCardWithEntities(entityId, entityToAdd.displayName, entityToAdd.childEntities);
-            }
+            view.add(_entities.get(entityId));
           } else {
             TheLogger.log("Warning", "Unknown entity inside view: $entityId");
           }
diff --git a/lib/view_class.dart b/lib/view_class.dart
index b5f949c..244bff5 100644
--- a/lib/view_class.dart
+++ b/lib/view_class.dart
@@ -15,25 +15,39 @@ class View {
     badges = {};
   }
 
-  void addBadge(String entityId) {
+  void add(Entity entity) {
+    if (!entity.isGroup) {
+      _addEntityWithoutGroup(entity);
+    } else {
+      _addCardWithEntities(entity);
+    }
+  }
+
+  void _addBadge(String entityId) {
     badges.addAll({entityId: Badge(entityId)});
   }
 
-  void addEntityWithoutGroup(Entity entity) {
-    String groupIdToAdd = "${entity.domain}.${entity.domain}$_count";
-    if (cards[groupIdToAdd] == null) {
-      addCard(groupIdToAdd, entity.domain);
+  void _addEntityWithoutGroup(Entity entity) {
+    if (UIBuilder.isBadge(entity.domain)) {
+      //This is badge
+      _addBadge(entity.entityId);
+    } else {
+      //This is a standalone entity
+      String groupIdToAdd = "${entity.domain}.${entity.domain}$_count";
+      if (cards[groupIdToAdd] == null) {
+        _addCard(groupIdToAdd, entity.domain);
+      }
+      cards[groupIdToAdd].addEntity(entity.entityId);
     }
-    cards[groupIdToAdd].addEntity(entity.entityId);
   }
 
-  void addCard(String entityId, String friendlyName) {
+  void _addCard(String entityId, String friendlyName) {
     cards.addAll({"$entityId": HACard(entityId, friendlyName)});
   }
 
-  void addCardWithEntities(String entityId, String friendlyName, List entities) {
-    cards.addAll({"$entityId": HACard(entityId, friendlyName)});
-    cards[entityId].addEntities(entities);
+  void _addCardWithEntities(Entity entity) {
+    cards.addAll({"${entity.entityId}": HACard(entity.entityId, entity.displayName)});
+    cards[entity.entityId].addEntities(entity.childEntities);
   }
 
 }
\ No newline at end of file