From 9de382b856866818e005e9e1f3ad36b79ee6eb19 Mon Sep 17 00:00:00 2001
From: jan <jan@orange.mpxd.net>
Date: Thu, 12 Oct 2023 01:29:57 -0700
Subject: [PATCH] Fix major bugs in presort

---
 masque/utils/pack2d.py | 29 +++++++++++++++++++++--------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/masque/utils/pack2d.py b/masque/utils/pack2d.py
index 52c0949..75a7d59 100644
--- a/masque/utils/pack2d.py
+++ b/masque/utils/pack2d.py
@@ -26,7 +26,7 @@ def maxrects_bssf(
     rejected_inds = set()
 
     if presort:
-        rotated_sizes = numpy.sort(rect_sizes, axis=0)  # shortest side first
+        rotated_sizes = numpy.sort(rect_sizes, axis=1)  # shortest side first
         rect_order = numpy.lexsort(rotated_sizes.T)[::-1]  # Descending shortest side
         rect_sizes = rect_sizes[rect_order]
 
@@ -80,26 +80,34 @@ def maxrects_bssf(
         r_top[:, 1] = loc[1] + rect_size[1]
 
         regions = numpy.vstack((regions[~intersects], r_lft, r_bot, r_rgt, r_top))
+
+    if presort:
+        unsort_order = rect_order.argsort()
+        rect_locs = rect_locs[unsort_order]
+        rejected_inds = set(unsort_order[list(rejected_inds)])
+
     return rect_locs, rejected_inds
 
 
-def guillotine_bssf_sas(rect_sizes: numpy.ndarray,
-                        regions: numpy.ndarray,
-                        presort: bool = True,
-                        allow_rejects: bool = True,
-                        ) -> tuple[numpy.ndarray, set[int]]:
+def guillotine_bssf_sas(
+        rects: ArrayLike,
+        containers: ArrayLike,
+        presort: bool = True,
+        allow_rejects: bool = True,
+        ) -> tuple[NDArray[numpy.float64], set[int]]:
     """
     sizes should be Nx2
     regions should be Mx4 (xmin, ymin, xmax, ymax)
     #TODO: test me!
     # TODO add rectangle-merge?
     """
-    rect_sizes = numpy.array(rect_sizes)
+    regions = numpy.array(containers, copy=False, dtype=float)
+    rect_sizes = numpy.array(rects, copy=False, dtype=float)
     rect_locs = numpy.zeros_like(rect_sizes)
     rejected_inds = set()
 
     if presort:
-        rotated_sizes = numpy.sort(rect_sizes, axis=0)  # shortest side first
+        rotated_sizes = numpy.sort(rect_sizes, axis=1)  # shortest side first
         rect_order = numpy.lexsort(rotated_sizes.T)[::-1]  # Descending shortest side
         rect_sizes = rect_sizes[rect_order]
 
@@ -138,6 +146,11 @@ def guillotine_bssf_sas(rect_sizes: numpy.ndarray,
         regions = numpy.vstack((regions[:rr], regions[rr + 1:],
                                 new_region0, new_region1))
 
+    if presort:
+        unsort_order = rect_order.argsort()
+        rect_locs = rect_locs[unsort_order]
+        rejected_inds = set(unsort_order[list(rejected_inds)])
+
     return rect_locs, rejected_inds