import numpy as np
from scipy.optimize import linprog
# https://[Log in to view URL]

def minimum_transportation_price(suppliers, consumers, costs):

    costs = np.array(costs)
    costs_flattened = costs.flatten()

    # Number of suppliers and consumers
    num_suppliers = len(suppliers)
    num_consumers = len(consumers)

    # Define the coefficients for the linear programming problem
    c = costs_flattened

    # Define the equality constraints matrix (A_eq) and vector (b_eq)
    A_eq = []
    b_eq = []

    # Constraints for suppliers
    for i in range(num_suppliers):
        constraint = [0] * (num_suppliers * num_consumers)
        constraint[i * num_consumers : (i + 1) * num_consumers] = [1] * num_consumers
        A_eq.append(constraint)
        b_eq.append(suppliers[i])

    # Constraints for consumers
    for j in range(num_consumers):
        constraint = [0] * (num_suppliers * num_consumers)
        for i in range(num_suppliers):
            constraint[i * num_consumers + j] = 1
        A_eq.append(constraint)
        b_eq.append(consumers[j])
    print(c, A_eq, b_eq)
    # Solve the linear programming problem
    res = linprog(c, A_eq=A_eq, b_eq=b_eq, method='revised simplex')
    
    # Reshape the result back to the original matrix shape
    solution = np.array(res.x).reshape(costs.shape)
    print("solution :", solution)
    return res.fun


tests = [
    [[10, 7, 13], [6, 20, 4], [[4, 12, 3], [20, 1, 6], [7, 0, 5]], 43],
    [[8, 15, 21], [8, 36], [[9, 16], [7, 13], [25, 1]], 288],
    [[31, 16], [14, 17, 16], [[41, 18, 0], [4, 16, 37]], 358],
    [[10, 20, 20], [5, 25, 10, 10], [[2, 5, 3, 0], [3, 4, 1, 4], [2, 6, 5, 2]], 150],
    [
        [13, 44, 27, 39, 17],
        [28, 12, 30, 17, 19, 34],
        [
            [6, 6, 12, 8, 13, 13],
            [7, 20, 5, 16, 11, 16],
            [4, 6, 19, 0, 2, 18],
            [1, 16, 6, 11, 8, 11],
            [5, 6, 11, 1, 6, 14],
        ],
        759,
    ],
    [
        [113, 68, 154, 135, 71, 238],
        [218, 95, 466], 
        [
            [13, 76, 70],
            [18, 100, 23],
            [72, 11, 66],
            [84, 75, 14],
            [89, 53, 93],
            [20, 45, 51]
        ],
        25348
    ],
]

for suppliers, consumers, costs, ans in tests:
    res = minimum_transportation_price(suppliers, consumers, costs)
    print(res, ans)