From 7c56f431ee9ff6f1d5978fdd678b781c1c2696e4 Mon Sep 17 00:00:00 2001
From: Zalo <gcrivelli@sig-gis.com>
Date: Mon, 23 Dec 2024 10:20:46 -0500
Subject: [PATCH] Improved services.py algorithm

---
 src/planscape/impacts/serializers.py |  1 -
 src/planscape/impacts/services.py    | 58 ++++++++++------------------
 src/planscape/impacts/views.py       | 22 +++++------
 3 files changed, 31 insertions(+), 50 deletions(-)

diff --git a/src/planscape/impacts/serializers.py b/src/planscape/impacts/serializers.py
index 6a1ae035e..7864c5be4 100644
--- a/src/planscape/impacts/serializers.py
+++ b/src/planscape/impacts/serializers.py
@@ -324,7 +324,6 @@ class TreatmentResultsPlotSerializer(serializers.Serializer):
 
 
 class StandQuerySerializer(serializers.Serializer):
-    # Use PrimaryKeyRelatedField to ensure stand_id is valid PK from Stand model
     stand_id = serializers.PrimaryKeyRelatedField(
         queryset=Stand.objects.all(), required=True
     )
diff --git a/src/planscape/impacts/services.py b/src/planscape/impacts/services.py
index b3596c78b..f6cd21d82 100644
--- a/src/planscape/impacts/services.py
+++ b/src/planscape/impacts/services.py
@@ -1,3 +1,7 @@
+#Temporary endpoint imports
+from collections import defaultdict
+from impacts.models import TreatmentResult, ImpactVariable
+
 import logging
 import itertools
 import json
@@ -530,56 +534,34 @@ def calculate_project_area_deltas(
     return results
 
 
-from impacts.models import TreatmentResult, ImpactVariable
-
-
 def get_treatment_results_table_data(treatment_plan, stand_id):
     """
-    Retrieves treatment results for RATE_OF_SPREAD and FLAME_LENGTH for the given stand and treatment plan.
-    Return a list of dictionaries: [{"year": ..., "rate_of_spread": ..., "flame_length": ...}, ...]
+    Retrieves a list of dictionaries, with each dictionary representing a year and its variables.
     """
-    # Filter results for the given stand and plan, for the two variables we need
+    # Use single query to retrieve all results at once
     queryset = TreatmentResult.objects.filter(
         treatment_plan=treatment_plan,
         stand_id=stand_id,
-        variable__in=[ImpactVariable.RATE_OF_SPREAD, ImpactVariable.FLAME_LENGTH],
-    )
+    ).values("year", "variable", "value")
 
     # If no results, return an empty list
     if not queryset.exists():
         return []
 
-    # Extract distinct years
-    years = queryset.values_list("year", flat=True).distinct().order_by("year")
-
-    # Build table rows
-    table_data = []
-    for year in years:
-        # For each year, get the corresponding records
-        year_results = queryset.filter(year=year)
-
-        # Extract the rate_of_spread value
-        rate_of_spread_val = (
-            year_results.filter(variable=ImpactVariable.RATE_OF_SPREAD)
-            .values_list("value", flat=True)
-            .first()
-        )
+    # Creates dictionary of { year: { variable: value } }
+    data_map = defaultdict(dict)
 
-        # Extract the flame_length value
-        flame_length_val = (
-            year_results.filter(variable=ImpactVariable.FLAME_LENGTH)
-            .values_list("Value", flat=True)
-            .first()
-        )
+    for row in queryset:
+        year = row["year"]
+        var = row["variable"]
+        val = row["value"]
+        data_map[year][var] = val
 
-        # Add a dictionary -- table row -- for each avaialble year
-        table_data.append(
-            {
-                "year": year,
-                "rate_of_spread": rate_of_spread_val,
-                "flame_length": flame_length_val,
-            }
-        )
+    # Convert data_map into a sorted list by year
+    table_data = []
+    for year in sorted(data_map.keys()):
+        row = {"year": year}
+        row.update(data_map[year])
+        table_data.append(row)
 
-    # returns list of rows to view
     return table_data
diff --git a/src/planscape/impacts/views.py b/src/planscape/impacts/views.py
index 4972f54c5..ed1da0fd4 100644
--- a/src/planscape/impacts/views.py
+++ b/src/planscape/impacts/views.py
@@ -1,3 +1,12 @@
+ # Temporary endpoint Imports
+from rest_framework import status, response
+from rest_framework.decorators import action
+from drf_spectacular.utils import extend_schema
+from planscape.serializers import BaseErrorMessageSerializer
+from impacts.serializers import StandQuerySerializer
+from impacts.services import get_treatment_results_table_data
+from impacts.models import TreatmentResult, ImpactVariable
+
 import logging
 from rest_framework import mixins, viewsets, response, status
 from rest_framework.decorators import action
@@ -257,14 +266,7 @@ def plot(self, request, pk=None):
         )
         return Response(data=data_to_plot, status=status.HTTP_200_OK)
 
-    # Imports
-    from rest_framework import status, response
-    from rest_framework.decorators import action
-    from drf_spectacular.utils import extend_schema
-    from planscape.serializers import BaseErrorMessageSerializer
-    from impacts.serializers import StandQuerySerializer
-    from impacts.services import get_treatment_results_table_data
-    from impacts.models import TreatmentResult, ImpactVariable
+   
 
     @extend_schema(
         description="Retrieve treatment result information for a specific stand.",
@@ -297,9 +299,7 @@ def stand_treatment_results(self, request, pk=None):
 
         if not table_data:
             return response.Response(
-                {
-                    "detail": "No treatment results found for this stand in the given treatment plan."
-                },
+                {"detail": "No treatment results found."},
                 status=status.HTTP_404_NOT_FOUND,
             )