using StatsBase

student_names = "Triana Forment";

"""
	Hungarian_algorithm(mat)

Solves a linear balanced assignment problem.

Input:
	- 'mat`: the cost matrix

Outputs:
	- `cost`: the resulting cost value
	- `matrix solution`: the matrix with the selected agent-task as its original cost value and the rest zeros. 
"""
function Hungarian_algorithm(mat)
	num_rows = size(mat)[1]
	num_cols = size(mat)[1]
		 @assert num_rows == num_cols throw(DimensionMismatch("It is not balanced. Number of sources and sinks need to match"))
	
	matrix = copy(mat)
	#Step 1: subtract its internal minimum from every column and every row 
	matrix .-= minimum(matrix, dims=2) # from rows
	matrix .-= minimum(matrix, dims=1) # from columns
	
	
	num_zeros = 0
	
	# Step 2: mark rows and columns that have the 0 entries 
	pos = [] 
    while num_zeros < num_rows
        pos, rows, cols = mark_matrix(matrix)
		rows = reverse(rows)
        num_zeros = length(rows) + length(cols)
		# Step 3: if marked rows & columns is less than num_rows, go to step 4.
        if num_zeros < num_rows
			# Step 4: adjust the matrix
            matrix = adjust_matrix(matrix, rows, cols)
		end
	end
	# Calculate total cost and matrix solution
	cost = 0
	matrix_solution = zeros(Int64, size(mat)[1], size(mat)[2])
	for i in 1:num_rows
		cost += mat[pos[i][1], pos[i][2]]
		matrix_solution[pos[i][1], pos[i][2]] = mat[pos[i][1], pos[i][2]]
	end
	return cost, matrix_solution
end

Hungarian_algorithm(mat)

"""
	find_min_row(matrix_zero, zero_list) 

Inputs:
	- `matrix_zero`: boolean matrix
	- `zero_list`: array to store the coordinates of the zeros in the row with fewest zeros.
	
Outputs:
	- `min_row`: tuple with the first element the number of zeros in the row with fewest zeros and the index of that row as second element
	- `zero_list`: list of coordinates tuples of the zeros found
"""
function find_min_row(matrix_zero, zero_list)
	min_row = [Inf, -1]
	
	for row in 1:size(matrix_zero)[1]
		#if true in matrix_zero[row, :] # Needs a 0 in the row
			# If the number of zeros < than the last min stored
			if sum(matrix_zero[row,:]) > 0 && min_row[1] > sum(matrix_zero[row,:])
				#stores the number of zeros and the row index
				min_row = [sum(matrix_zero[row,:]), row]
			end
	end
	# Get the column index 
	zero_index = findall(x->x==true, matrix_zero[min_row[2],:])[1]  
	# Store tuple in zero_list the tuple of coordinates
	append!(zero_list, [(min_row[2], zero_index)]) 
	# Mark the specific row and column as false.
	matrix_zero[min_row[2], :] .= false
	matrix_zero[:, zero_index] .= false
	return min_row, zero_list
end

find_min_row(matrix_zero, zero_list)

"""
	mark_matrix(mat)

Perform steps 2 and 3: It marks the rows and columns that have the 0 entries such that the fewest lines possible are drawn. If there are "m" marked rows and columns, an optimal assignment solution can be found.  

Input:
	- `mat` : cost matrix already modified by step 1 or an adjusted matrix by `adjust_matrix`
	
Outputs:
	- `zero_list` : coordinates of the zeros found by `find_min_row`
	- `marked_rows`: array of marked rows 
	- `marked_cols`: array of marked columns
"""
function mark_matrix(mat)
	
	#Transform the matrix to boolean matrix(0 = true, others = false)
	cur_mat = mat
	boolean_mat = iszero.(mat)
	boolean_mat_copy = copy(boolean_mat)
	#Recording possible answer positions by zero_list
    zero_list = []	
    while true in boolean_mat_copy
        find_min_row(boolean_mat_copy, zero_list)
	end
    #Recording the row indexes
    zero_list_row = []
    for i in 1:length(zero_list)
        append!(zero_list_row, zero_list[i][1])
	end

	# Get non marked rows
    non_marked_row = (x->collect(x))(setdiff(Set(1:size(cur_mat)[1]), Set(zero_list_row)))
    
    marked_cols = []
    flag = 0 
    while flag == 0 # Enter in the loop until there's no more unmarked cols and rows
        flag = 1 
        for i in 1:length(non_marked_row)
            row_array = boolean_mat[non_marked_row[i], :]
            for j in 1:length(row_array)
				# Find unmarked 0 elements in the corresponding column
				if row_array[j] == true && j ∉ marked_cols
                    # Store column index in "marked_cols"
                    append!(marked_cols, j)
                    flag = 0
				end
			end
		end
		zero_rows = [t[1] for t in zero_list]
		zero_cols = [t[2] for t in zero_list]
		for i in 1:length(zero_list)
			# If in "zero_list" coordinates there was a 0 marked in cols,
			# but not in marked_rows, add its row index to "non_marked_row"
            if zero_rows[i] ∉ non_marked_row && zero_cols[i] in marked_cols
                append!(non_marked_row, zero_rows[i])
                flag = 0
			end
		end
	end
	# Add those the indexes not stored in "non_marked_row" to "marked_rows"
	marked_rows = (x->collect(x))(setdiff(Set(1:size(cur_mat)[1]), Set(non_marked_row)))   
    return zero_list, marked_rows, marked_cols
end			

mark_matrix(mat)

"""
	adjust_matrix(mat, marked_rows, marked_cols)

Perform step 4: Adjust the matrix by finding the smallest entry not marked by any marked row and by any column. Subtract this entry from each row that isn’t marked, and then add it to each column that is marked. 

Inputs:
	- `mat` : boolean matrix created and modified by `mark_matrix`
	- `marked_rows`: array of marked rows produced by `mark_matrix`
	- `marked_cols`: array of marked columns produced by `mark_matrix`

Output:
	- `adjusted`: the adjusted matrix
"""
function adjust_matrix(mat, marked_rows, marked_cols)
	adjusted = mat
    min_values = []
    n = size(adjusted)[1]
	
    # Find the minimum value for that is not in marked_rows and not in marked_cols
    for row in 1:n
        if row ∉ marked_rows
            for i in 1:n
                if i ∉ marked_cols
                    append!(min_values, adjusted[row,i])
				end
			end
		end
	end
    min_value = minimum(min_values)
	# Substract that min value to all elements not in marked rows and columns
    for row in 1:n
        if row ∉ marked_rows
            for i in 1:n
                if i ∉ marked_cols
                    adjusted[row, i] = adjusted[row, i] - min_value
				end
			end
		end
	end
    # Add the min value to elements that are in both marked row and columns
    for row in marked_rows
        for col in marked_cols
            adjusted[row, col] = adjusted[row,col] + min_value
		end
	end
    return adjusted
end

adjust_matrix(mat, marked_rows, marked_cols)

cost_matrix = [8 6 2 9 2;
               1 2 5 3 9;
               5 3 9 6 5;
               9 2 5 8 7;
               2 5 3 6 1]

5×5 Matrix{Int64}:
 8  6  2  9  2
 1  2  5  3  9
 5  3  9  6  5
 9  2  5  8  7
 2  5  3  6  1

Hungarian_algorithm(cost_matrix)

(12, 5×5 Matrix{Int64}:
 0  0  2  0  0
 1  0  0  0  0
 0  0  0  6  0
 0  2  0  0  0
 0  0  0  0  1)

begin 
	benefit_matrix = copy(cost_matrix)
	benefit_matrix .= - cost_matrix 
end

5×5 Matrix{Int64}:
 -8  -6  -2  -9  -2
 -1  -2  -5  -3  -9
 -5  -3  -9  -6  -5
 -9  -2  -5  -8  -7
 -2  -5  -3  -6  -1

Hungarian_algorithm(benefit_matrix)

(-41, 5×5 Matrix{Int64}:
  0   0   0  -9   0
  0   0   0   0  -9
  0   0  -9   0   0
 -9   0   0   0   0
  0  -5   0   0   0)