Updated data extensder

iq_mini_4489/src/data_extender.py

@@ -77,11 +77,34 @@ def main():
     # Combine original data with new rotated entries
     extended_data = data + new_entries
     
-    # Save the extended data to a new JSON file
+    # Define the arrays for pinA, pinB, and pinC containing tuples of coordinates
+    pinA = [(0, 0), (1, 0), (0, 1), (1, 1), (2, 1), (3, 1), (1, 2), (1, 3)]  # x pins
+    pinB = [(3, 0), (4, 0), (4, 1), (4, 2), (4, 3), (3, 4), (4, 4)]          # y pins
+    pinC = [(2, 2), (3, 2), (0, 3), (2, 3), (0, 4), (1, 4), (2, 4)]          # z pins
+
+    # Define the array for forbidden positions
+    forbidden_positions = [(2, 0), (0, 2), (3, 3)]  # o positions
+
+    # Remove forbidden positions from the extended data
+    filtered_data = []
+    for entry in extended_data:
+        pinA, pinB, pinC = entry['pinA'], entry['pinB'], entry['pinC']
+        if pinA not in forbidden_positions and pinB not in forbidden_positions and pinC not in forbidden_positions:
+            filtered_data.append(entry)
+    
     with open('data_extended.json', 'w') as f:
         json.dump(extended_data, f, indent=4)
     
-    print("Extended data saved to 'data_extended.json'.")
-
+    with open("data_filtered.json", "w") as f:
+        json.dump(filtered_data, f, indent=4)
+        
+    # Print distribution of solutions in the filtered data
+    solution_counts = [entry['solution_count'] for entry in filtered_data]
+    unique_solution_counts = set(solution_counts)
+    print("Distribution of solutions in the filtered data:")
+    for count in unique_solution_counts:
+        count_entries = [entry for entry in filtered_data if entry['solution_count'] == count]
+        print(f"Solutions: {count}, Count: {len(count_entries)}")
+    
 if __name__ == "__main__":
     main()