From 6b65ac77b38d5d6d7ec0316497c4112569869e3c Mon Sep 17 00:00:00 2001
From: Dana Lambert <danalambert@catalyst.net.nz>
Date: Thu, 4 Nov 2021 11:18:42 +1300
Subject: [PATCH] Add command to initially load habitat and zone data

---
 README.md                                     |  1 +
 .../api/management/commands/loadsitedata.py   | 62 +++++++++++++++++++
 dev                                           | 11 ++++
 3 files changed, 74 insertions(+)
 create mode 100644 backend/right_tree/api/management/commands/loadsitedata.py

diff --git a/README.md b/README.md
index f69e28f..795fa2b 100644
--- a/README.md
+++ b/README.md
@@ -108,6 +108,7 @@ A summary of available commands are outlined below. Note that if the command req
 | `reset_plants` | Performs the custom `resetplants` command in the backend container. This removes all plant entries from the database. | Yes
 | `load_plant_fixtures` | Loads the `/backend/right_tree/api/data/fixtures/plants.json` fixture. Requires the `plants.json` file to be created (`./dev create_plant_fixtures`) and the plant table to be empty (`./dev reset_plants`). | Yes
 | `load_plants` | Creates plants fixtures and loads them into a fresh plant table in the database. Requires the fixtures to be applied and shapefiles loaded. | Yes
+| `load_sites_from_spreadsheet` | Loads site spreadsheet data the database initially (replaced with fixtures containing further information) | Yes
 | `populate_database` | Populates the `right_tree` database with base data (fixtures), provided shapefiles and plant spreadsheet data. Requires the database to be created. | No
 | `init_database` | Creates and populates the database | No
 | `reset_database` | Removes, recreates and populates the database | No
diff --git a/backend/right_tree/api/management/commands/loadsitedata.py b/backend/right_tree/api/management/commands/loadsitedata.py
new file mode 100644
index 0000000..a714db9
--- /dev/null
+++ b/backend/right_tree/api/management/commands/loadsitedata.py
@@ -0,0 +1,62 @@
+from django.core.management.base import BaseCommand
+
+from pathlib import Path
+
+import right_tree.api.data
+from ._spreadsheet_helpers import *
+from right_tree.api.models import Habitat, Zone
+
+SPREADSHEET_FILENAME = 'plant_data.xlsx'
+SITE_DATA_START_COL = 25
+HABITAT_ROW = 1
+ZONE_NAME_ROW = 2
+ZONE_VARIANT_ROW = 3
+ZONE_REFINED_VARIANT_ROW = 5
+
+DATA_DIR_PATH = Path(right_tree.api.data.__file__).resolve().parent
+SPREADSHEET = get_spreadsheet(DATA_DIR_PATH, SPREADSHEET_FILENAME)
+
+
+def load_habitat_zone_data(sheet):
+    """ Loads habitat and zone objects from data defined in the spreadsheet.
+    """
+
+    habitats = set()
+    zones = set()
+
+    current_habitat = current_zone = current_variant = current_refined_variant = None
+
+    for col in sheet.iter_cols(min_col=SITE_DATA_START_COL, min_row=ZONE_NAME_ROW, max_row=ZONE_REFINED_VARIANT_ROW, values_only=True):
+        habitat, zone_name, zone_variant, zone_refined_variant = col
+
+        if habitat is not None:
+            current_habitat = habitat
+            current_zone = current_variant = current_refined_variant = None
+
+        current_zone = zone_name if zone_name is not None else current_zone
+        current_variant = zone_variant if zone_variant is not None else current_variant
+        current_refined_variant = zone_refined_variant if zone_refined_variant is not None else current_refined_variant
+
+        habitats.add(current_habitat)
+        zones.add((current_zone, current_variant, current_refined_variant))
+
+    print("Creating Habitat Objects")
+    print("--------------------------")
+    for habitat in habitats:
+        habitat_obj = Habitat(name=habitat)
+        habitat_obj.save()
+
+    print("Creating Zone Objects")
+    print("--------------------------")
+    for zone, variant, refined_variant in zones:
+        zone_obj = Zone(name=zone, variant=variant,
+                        refined_variant=refined_variant)
+        zone_obj.save()
+
+
+class Command(BaseCommand):
+    help = 'Ingests the site spreadsheet data into the database'
+
+    def handle(self, *args, **options):
+        self.stdout.write('Creating habitat/zone fixtures...')
+        load_habitat_zone_data(SPREADSHEET)
diff --git a/dev b/dev
index 156b514..dd892a7 100755
--- a/dev
+++ b/dev
@@ -1,5 +1,11 @@
 #!/bin/bash
 
+# Load .env file if it exists
+if [ -f .env ]
+then
+  export $(cat .env | sed 's/#.*//g' | xargs)
+fi
+
 cmd_create_database() {
     echo "Creating right_tree database..."
     docker-compose down --remove-orphans --volumes
@@ -53,6 +59,11 @@ cmd_load_plants() {
     cmd_load_plant_fixtures    
 }
 
+cmd_load_sites_from_spreadsheet() {
+    echo "Loading habitats and zones..."
+    docker-compose exec django-backend python manage.py loadsitedata
+}
+
 cmd_populate_database() {
     echo "Populating the database..."
     docker-compose up -d django-backend postgres