Implement discount budget implementation

lib/planet_express_express.rb

@@ -4,6 +4,7 @@ require_relative "planet_express_express/version"
 require_relative "planet_express_express/constants"
 require_relative "planet_express_express/shipment_reader"
 require_relative "planet_express_express/rules"
+require_relative "planet_express_express/discount_budget"
 
 module PlanetExpressExpress
   class Error < StandardError; end

lib/planet_express_express/discount_budget.rb

@@ -0,0 +1,64 @@
+module PlanetExpressExpress
+  class DiscountBudget
+    def initialize(limit_per_month)
+      @limit_per_month = limit_per_month
+    end
+
+    def limit_discounts(shipments)
+      keys = shipments.keys
+      positions = Array.new(keys.length, 0)
+
+      # We expect all keys to have the same length
+      total_length = shipments[keys.first].length
+
+      for month in 0..(total_length - 1) do
+        reset_positions(positions)
+        budget = @limit_per_month
+
+        while true
+          # Find the next shipment
+          min_index = Float::INFINITY
+          next_val = nil # key_index
+
+          for key_index in 0..(keys.length - 1) do
+            current_position = positions[key_index]
+
+            shipment = shipments[keys[key_index]][month][current_position]
+            next if shipment.nil?
+
+            if shipment[:id] < min_index
+              next_val = key_index
+              min_index = shipment[:id]
+            end
+          end
+
+          break if next_val.nil?
+          # End find the next shipment
+
+          shipment = shipments[keys[next_val]][month][positions[next_val]]
+          positions[next_val] = positions[next_val] + 1
+
+          budget = apply_discount(budget, shipment)
+        end
+      end
+    end
+
+    private
+
+    attr_reader :limit_per_month
+
+    def apply_discount(available_budget, shipment)
+      if shipment[:discount] > available_budget
+        shipment[:discount] = available_budget
+        return 0
+      end
+      available_budget - shipment[:discount]
+    end
+
+    def reset_positions(positions)
+      for i in 0..(positions.length - 1) do
+        positions[i] = 0
+      end
+    end
+  end
+end