optimize load balance for inventory list in pos - add new fields to stake holdes

apps/herd/services/services.py

@@ -3,36 +3,27 @@ from apps.herd.models import Rancher
 from apps.livestock.models import LiveStock
 from apps.warehouse.models import InventoryEntry
 from apps.product.models import Quota
+from django.db.models import Count, Q
 import typing
 
 
 def get_rancher_statistics(rancher: Rancher = None) -> typing.Any:
     """ get statistics of a rancher """  # noqa
 
-    herds = rancher.herd.all()  # noqa
-    herd_count = herds.count()
+    livestocks = LiveStock.objects.filter(herd__rancher=rancher)  # noqa
 
-    livestocks = LiveStock.objects.filter(herd__in=herds)  # noqa
+    stats = livestocks.aggregate(
+        herd_count=Count("herd", distinct=True),
+        light_count=Count('id', filter=Q(weight_type='L')),
+        heavy_count=Count('id', filter=Q(weight_type='H')),
+        sheep_count=Count('id', filter=Q(type__name='گوسفند')),  # noqa
+        goat_count=Count('id', filter=Q(type__name='بز')),
+        cow_count=Count('id', filter=Q(type__name='گاو')),
+        camel_count=Count('id', filter=Q(type__name='شتر')),
+        horse_count=Count('id', filter=Q(type__name='بز')),
+    )
 
-    light_count = livestocks.filter(weight_type='L').count()
-    heavy_count = livestocks.filter(weight_type='H').count()
-
-    sheep_count = livestocks.filter(type__name="گوسفند").count()  # noqa
-    goat_count = livestocks.filter(type__name="بز").count()
-    cow_count = livestocks.filter(type__name="گاو").count()
-    camel_count = livestocks.filter(type__name="شتر").count()
-    horse_count = livestocks.filter(type__name="اسب").count()
-
-    return {
-        "herd_count": herd_count,
-        "light_count": light_count,
-        "heavy_count": heavy_count,
-        "sheep_count": sheep_count,
-        "goat_count": goat_count,
-        "cow_count": cow_count,
-        "camel_count": camel_count,
-        "horse_count": horse_count,
-    }
+    return stats
 
 
 def rancher_quota_weight(rancher, inventory_entry: InventoryEntry):
@@ -51,60 +42,38 @@ def rancher_quota_weight(rancher, inventory_entry: InventoryEntry):
     }
 
     quota: Quota = inventory_entry.distribution.quota
-    allocations = quota.livestock_allocations.all()  # list of quota live stock allocations
+    # list of quota live stock allocations
+    allocations = list(quota.livestock_allocations.all().select_related('livestock_type'))
+    # list of quota incentive plans
+    incentive_plans = list(quota.incentive_assignments.all().select_related('livestock_type'))
 
     livestock_counts = get_rancher_statistics(rancher)
 
     total_weight = 0
-    alloc_details = {}
-    plan_details = {}
-    result_list = []
+    merged = {}
 
-    # list of quota allocations, get allocations weight on any animal type
-    for alloc in allocations:  # noqa
-        if alloc.livestock_type:
-            animal_type = alloc.livestock_type.name
-            per_head = alloc.quantity_kg
+    for item in allocations + incentive_plans:  # noqa
+        if item.livestock_type:
+            animal_type = item.livestock_type.name
+            per_head = item.quantity_kg
             count = livestock_counts.get(live_stock_meta.get(animal_type), 0)
 
             weight = per_head * count
-            alloc_details[animal_type] = {"weight": weight, "type": alloc.livestock_type.weight_type}
             total_weight += weight
 
-    # list of quota incentive plans, get plans weight on any animal type
-    incentive_plans = quota.incentive_assignments.all()
-    for plan in incentive_plans:  # noqa
-        if plan.livestock_type:
-            animal_type = plan.livestock_type.name
-            per_head = plan.quantity_kg
-            count = livestock_counts.get(live_stock_meta.get(animal_type), 0)
+            if animal_type not in merged:
+                merged[animal_type] = {
+                    "weight": weight,
+                    "type": item.livestock_type.weight_type
+                }
+            else:
+                merged[animal_type]['weight'] += weight
 
-            weight = per_head * count
-            plan_details[animal_type] = {"weight": weight, "type": plan.livestock_type.weight_type}
-            total_weight += weight
-
-    # summation of incentive plans & livestock allocations animal types weight
-    result_details = {"total": total_weight, 'by_type': {}}
-    all_keys = set(alloc_details.keys()) | set(plan_details.keys())  # get all keys from plan & allocations data
-
-    for key in all_keys:
-        alloc_weight = alloc_details.get(key, {}).get("weight", 0)  # total weight of quota livestock allocations data
-        plan_weight = plan_details.get(key, {}).get("weight", 0)  # total weight of quota incentive plan data
-
-        # get animal type (Heavy, Light)
-        animal_type = alloc_details.get(
-            key, {}
-        ).get("type") or plan_details.get(
-            key, {}
-        ).get("type")
-
-        # final result, total weights
-        result_list.append({
+    return {
+        "total_weight": total_weight,
+        "by_type": [{
             "name": key,
-            "weight": alloc_weight + plan_weight,
-            "type": animal_type
-        })
-
-    result_details['by_type'] = result_list
-
-    return result_details
+            "weight": value['weight'],
+            "type": value['type']
+    }for key, value in merged.items()]
+    }