import math

MAX_SAME_GROUP = 2
MAX_PRODUCT_QUANTITY = 3
MAX_DIFFICULT_QUANTITY = 2
MAX_FRIDGE_QUANTITY = 6
MAKE_ALONE_PRODUCTS = ["caipirinha", "sangria_2p", "sangria_3p", "sangria_4p", "sangria_5p", "aperol", "espresso_martini", "long_island_iced_tea", "guiness_shot"]
MAX_DIFFICULT_PRODUCTS = ["caipirinha", "sangria_2p", "sangria_3p", "sangria_4p", "sangria_5p", "espresso_martini", "long_island_iced_tea"]

def select_waiter_with_least_load(waiter_load, waiter_products, waiter_order_load, exclude_waiter=None):
    #sort waiters, minimum load first
    sorted_waiters = sorted(waiter_load.items(), key=lambda x: x[1])
    for waiter, load in sorted_waiters:
        if waiter_products[waiter] == {} and waiter != exclude_waiter:
            return waiter
        else:
            continue
    #return waiter with least load in the order
    sorted_waiters_order = sorted(waiter_order_load.items(), key=lambda x: x[1])
    return sorted_waiters_order[0][0]
            

class OrderDistributor:
    def __init__(self, product_costs, waiters):
        self.product_costs = product_costs
        self.groups = list(product_costs.keys())
        self.flat_product_costs = {}
        for group in self.product_costs:
            for product in self.product_costs[group]:
                self.flat_product_costs[product] = self.product_costs[group][product]
        self.waiters = waiters
        self.waiter_load = {w: 0 for w in self.waiters}

    def distribute_single_order(self, order):
        def assign_to_waiter(waiter, product, assign_qty):
            waiter_order_load[waiter] += self.flat_product_costs[product] * assign_qty
            waiter_quantity[waiter] += assign_qty
            self.waiter_load[waiter] += self.flat_product_costs[product] * assign_qty
            if product in waiter_products[waiter]:
                waiter_products[waiter][product] += assign_qty
            else:
                waiter_products[waiter][product] = assign_qty


        waiter_products = {w: {} for w in self.waiters}
        waiter_order_load = {w: 0 for w in self.waiters}    

        for group in self.groups:
            waiter = select_waiter_with_least_load(self.waiter_load, waiter_products, waiter_order_load)
            waiter_quantity = {w: 0 for w in self.waiters}
            for product in order:
                # Skip products not in this group
                if product not in self.product_costs[group]:
                    continue

                # if product in group:
                quantity = order[product]
                group_limit = MAX_PRODUCT_QUANTITY
                if group == "fridges":
                    group_limit = MAX_FRIDGE_QUANTITY
                if product in MAKE_ALONE_PRODUCTS:
                    group_limit = MAX_DIFFICULT_QUANTITY

                if waiter_quantity[waiter] + quantity >= group_limit:
                    assign_qty = math.floor(quantity / 2)
                    assign_to_waiter(waiter, product, assign_qty)
                    remaining_qty = quantity - assign_qty
                    if remaining_qty > 0:
                        # Find next waiter with lowest load who isn't the current one
                        while True:
                            next_waiter = select_waiter_with_least_load(self.waiter_load, waiter_products, waiter_order_load, exclude_waiter=waiter)
                            if next_waiter != waiter:
                                break
                        assign_to_waiter(next_waiter, product, remaining_qty)
                else:
                    assign_to_waiter(waiter, product, quantity)
                
                if group == "alone_make_products" or product in MAX_DIFFICULT_PRODUCTS or waiter_products[waiter][product] >= 2:
                    waiter = select_waiter_with_least_load(self.waiter_load, waiter_products, waiter_order_load, exclude_waiter=waiter)
            
                
                

        # Build the distribution result for this order
        distribution = {}
        for product in order:
            for waiter in waiter_products:
                if product in waiter_products[waiter]:
                    if product not in distribution:
                        distribution[product] = {}
                    distribution[product][waiter] = waiter_products[waiter][product]
        print(self.waiter_load)
        return distribution, 
    
    def distribute_orders(self, orders):
        self.waiter_load = {w: 0 for w in self.waiters}
        final_order_list = []
        for order in orders:
            final_order_list.append(self.distribute_single_order(order))
        return final_order_list, self.waiter_load