diff --git a/client/app.js b/client/app.js
index c94b8cb2..228c1c07 100644
--- a/client/app.js
+++ b/client/app.js
@@ -31,5 +31,5 @@ app.get('/demo/units', (req, res) => demoDataGenerator.units(req, res));
 app.get('/demo/logs', (req, res) => demoDataGenerator.logs(req, res));
 app.get('/demo/bullseyes', (req, res) => demoDataGenerator.bullseyes(req, res));
 app.get('/demo/airbases', (req, res) => demoDataGenerator.airbases(req, res));
-
+app.get('/demo/mission', (req, res) => demoDataGenerator.mission(req, res));
 
diff --git a/client/demo.js b/client/demo.js
index 03e08ef8..f6e70fe0 100644
--- a/client/demo.js
+++ b/client/demo.js
@@ -588,6 +588,11 @@ class DemoDataGenerator {
         }};
         res.send(JSON.stringify(ret));
     };
+
+    mission(req, res){
+        var ret = {theatre: "NTTR"};
+        res.send(JSON.stringify(ret));
+    }
     
     generateRandomUnitsDemoData(unitsNumber) 
     {
diff --git a/client/src/index.ts b/client/src/index.ts
index a16d03a5..d4c83c8f 100644
--- a/client/src/index.ts
+++ b/client/src/index.ts
@@ -182,7 +182,7 @@ function setupEvents() {
     });
 
     document.addEventListener("toggleUnitVisibility", (ev: CustomEventInit) => {
-        document.body.toggleAttribute("data-hide-" + ev.detail.unitType);
+        document.body.toggleAttribute("data-hide-" + ev.detail.category);
     });
 
     document.addEventListener( "closeDialog", (ev: CustomEventInit) => {
diff --git a/client/src/units/unit.ts b/client/src/units/unit.ts
index 71c348e7..7c32f67e 100644
--- a/client/src/units/unit.ts
+++ b/client/src/units/unit.ts
@@ -59,6 +59,7 @@ export class Unit extends Marker {
 
     #selectable: boolean;
     #selected: boolean = false;
+    #hidden: boolean = false;
 
     #preventClick: boolean = false;
 
@@ -100,6 +101,15 @@ export class Unit extends Marker {
         this.#pathPolyline.addTo(getMap());
         this.#targetsPolylines = [];
 
+        // Deselect units if they are hidden
+        document.addEventListener("toggleCoalitionVisibility", (ev: CustomEventInit) => {
+            setTimeout(() => {this.setSelected(this.getSelected() && !this.getHidden())}, 300);
+        });
+    
+        document.addEventListener("toggleUnitVisibility", (ev: CustomEventInit) => {
+            setTimeout(() => {this.setSelected(this.getSelected() && !this.getHidden())}, 300);
+        });
+
         this.setData(data);
     }
 
@@ -111,7 +121,6 @@ export class Unit extends Marker {
             this.getBaseData().alive != data.baseData.alive || this.#forceUpdate || !getMap().hasLayer(this))
             updateMarker = true;
             
-
         if (data.baseData != undefined)
         {
             for (let key in this.#data.baseData)
@@ -161,7 +170,7 @@ export class Unit extends Marker {
         }
 
         /* Dead units can't be selected */
-        this.setSelected(this.getSelected() && this.getBaseData().alive)
+        this.setSelected(this.getSelected() && this.getBaseData().alive && !this.getHidden())
 
         if (updateMarker)
             this.#updateMarker();
@@ -205,9 +214,9 @@ export class Unit extends Marker {
 
     setSelected(selected: boolean) {
         /* Only alive units can be selected. Some units are not selectable (weapons) */
-        if ((this.getBaseData().alive || !selected) && this.#selectable && this.#selected != selected) {
+        if ((this.getBaseData().alive || !selected) && this.getSelectable() && this.getSelected() != selected) {
             this.#selected = selected;
-            this.getElement()?.querySelector( `[data-object|="unit"]` )?.toggleAttribute( "data-is-selected" );
+            this.getElement()?.querySelector(`[data-object|="unit"]`)?.toggleAttribute("data-is-selected");
             if (selected)
                 document.dispatchEvent(new CustomEvent("unitSelection", { detail: this }));
             else
@@ -243,8 +252,13 @@ export class Unit extends Marker {
         this.getTaskData().activePath = undefined;
     }
 
+    setHidden(hidden: boolean)
+    {
+        this.#hidden = hidden;
+    }
+
     getHidden() {
-        return false;
+        return (<HTMLElement>this.getElement()?.querySelector(`.unit`))?.offsetParent === null; 
     }
 
     getLeader() {
@@ -474,9 +488,7 @@ export class Unit extends Marker {
 }
 
 export class AirUnit extends Unit {
-    getHidden() {
-        return false;
-    }
+
 }
 
 export class Aircraft extends AirUnit {
@@ -533,10 +545,6 @@ export class GroundUnit extends Unit {
             </div>
         `);
     }
-
-    getHidden() {
-        return false;
-    }
 }
 
 export class NavyUnit extends Unit {
@@ -549,10 +557,6 @@ export class NavyUnit extends Unit {
             </div>
         `);
     }
-
-    getHidden() {
-        return false;
-    }
 }
 
 export class Weapon extends Unit {
diff --git a/client/views/navbar.ejs b/client/views/navbar.ejs
index 47a2f3c3..11b8e83c 100644
--- a/client/views/navbar.ejs
+++ b/client/views/navbar.ejs
@@ -28,11 +28,11 @@
     </div>
     
     <div id="unit-visibility-control" class="ol-group">
-        <button id="unit-visibility-control-aircraft" title="Toggle aircraft' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "unitType": "aircraft" }'></button>
-        <button id="unit-visibility-control-ground" title="Toggle ground units' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "unitType": "ground" }'></button>
-        <button id="unit-visibility-control-sam" title="Toggle SAMs' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "unitType": "sam" }'></button>
-        <button id="unit-visibility-control-threat" title="Toggle threat rings' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "unitType": "threat" }'></button>
-        <button id="unit-visibility-control-naval" title="Toggle naval visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "unitType": "naval" }'></button>
+        <button id="unit-visibility-control-aircraft" title="Toggle aircraft' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "category": "aircraft" }'></button>
+        <button id="unit-visibility-control-ground" title="Toggle ground units' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "category": "ground" }'></button>
+        <button id="unit-visibility-control-sam" title="Toggle SAMs' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "category": "sam" }'></button>
+        <button id="unit-visibility-control-threat" title="Toggle threat rings' visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "category": "threat" }'></button>
+        <button id="unit-visibility-control-naval" title="Toggle naval visibility" data-on-click="toggleUnitVisibility" data-on-click-params='{ "category": "naval" }'></button>
     </div>
 
     <div id="coalition-visibility-control" class="ol-group ol-group-button-toggle">