
# -*- coding: utf-8 -*-


from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function


''' lapjv.py - Jonker-Volgenant algorithm for linear assignment problem.

This is an implementation of the Jonker-Volgenant algorithm for solving
the linear assignment problem. The code is derived from the paper,

R. Jonker, A. Volgenant, "A Shortest Augmenting Path Algorithm for Dense
and Sparse Linear Assignment Problems", Computing 38, p 325-340, 1987

CellProfiler is distributed under the GNU General Public License,
but this file is licensed under the more permissive BSD license.
See the accompanying file LICENSE for details.

Copyright (c) 2003-2009 Massachusetts Institute of Technology
Copyright (c) 2009-2013 Broad Institute
All rights reserved.
'''

import numpy as np

try:  # pragma: no cover
    from ._lapjv import reduction_transfer
    from ._lapjv import augmenting_row_reduction
    from ._lapjv import augment
except ImportError:
    # Try on the fly Cython compilation
    import pyximport
    pyximport.install(setup_args={'include_dirs': [np.get_include()]})

    from ._lapjv import reduction_transfer
    from ._lapjv import augmenting_row_reduction
    from ._lapjv import augment

__all__ = []


def lapjv(i, j, costs, wants_dual_variables=False,
          augmenting_row_reductions=2, use_slow=False):  # pragma: no cover
    '''Sparse linear assignment solution using Jonker-Volgenant algorithm

    i,j - similarly-sized vectors that pair the object at index i[n] with
          the object at index j[n]

    costs - a vector of similar size to i and j that is the cost of pairing
            i[n] with j[n].

    wants_dual_variables - the dual problem reduces the costs using two
            vectors, u[i] and v[j] where the solution is the maximum value of
            np.sum(u) + np.sum(v) where cost[i,j] - u[i] - v[j] >=  0.
            Set wants_dual_variables to True to have u and v returned in
            addition to the assignments.

    augmenting_row_reductions - the authors suggest that augmenting
            row reduction be performed twice to optimize the u and v
            before the augmenting stage. The caller can choose a
            different number of reductions by supplying a different
            value here.
    use_slow: Bool, default False:
            use the pure python implementation, useful for debugging purpose

    All costs not appearing in i,j are taken as infinite. Each i in the range,
    0 to max(i) must appear at least once and similarly for j.

    returns (x, y), the pairs of assignments that represent the solution
    or (x, y, u, v) if the dual variables are requested.
    '''
    i = np.atleast_1d(i).astype(int)
    j = np.atleast_1d(j).astype(int)
    costs = np.atleast_1d(costs)

    assert len(i) == len(j), "i and j must be the same length"
    assert len(i) == len(costs), "costs must be the same length as i"

    #
    # Find the number of i with non-infinite cost for each j
    #
    j_count = np.bincount(j)
    assert not np.any(j_count == 0), "all j must be paired with at least one i"
    #
    # if you order anything by j, this is an index to the minimum for each j
    #
    j_index = np.hstack([[0], np.cumsum(j_count[:-1])])
    #
    # Likewise for i
    #
    i_count = np.bincount(i)
    assert not np.any(i_count == 0), "all i must be paired with at least one j"
    i_index = np.hstack([[0], np.cumsum(i_count[:-1])])

    n = len(j_count)  # dimension of the square cost matrix
    assert n == len(i_count), "There must be the same number of unique i and j"

    # # # # # # # #
    #
    # Variable initialization:
    #
    # The output variables:
    #
    # x - for each i, the assigned j. -1 indicates uninitialized
    # y - for each j, the assigned i
    # u, v - the dual variables
    #
    # A value of x = n or y = n means "unassigned"
    #
    x = np.ascontiguousarray(np.ones(n, np.uint32) * n)
    y = np.ascontiguousarray(np.ones(n, np.uint32) * n, np.uint32)
    u = np.ascontiguousarray(np.zeros(n, np.float64))

    # # # # # # # #
    #
    # Column reduction
    #
    # # # # # # # #
    #
    # For a given j, find the i with the minimum cost.
    #
    order = np.lexsort((-i, costs, j))
    min_idx = order[j_index]
    min_i = i[min_idx]
    #
    # v[j] is assigned to the minimum cost over all i
    #
    v = np.ascontiguousarray(costs[min_idx], np.float64)
    #
    # Find the last j for which i was min_i.
    #
    x[min_i] = np.arange(n).astype(np.uint32)
    y[x[x != n]] = np.arange(n).astype(np.uint32)[x != n]
    #
    # Three cases for i:
    #
    # i is not the minimum of any j - i goes on free list
    # i is the minimum of one j - v[j] remains the same and y[x[j]] = i
    # i is the minimum of more than one j, perform reduction transfer
    #
    assignment_count = np.bincount(min_i[min_i != n])
    # print assignment_count
    assignment_count = np.hstack(
        (assignment_count, np.zeros(n - len(assignment_count), int)))
    free_i = assignment_count == 0
    one_i = assignment_count == 1
    # print one_i
    # order = np.lexsort((costs, i)) Replace with this after all is done
    order = np.lexsort((j, i))
    j = np.ascontiguousarray(j[order], np.uint32)
    costs = np.ascontiguousarray(costs[order], np.float64)
    i_index = np.ascontiguousarray(i_index, np.uint32)
    i_count = np.ascontiguousarray(i_count, np.uint32)
    if use_slow:
        if np.any(one_i):
            print(np.ascontiguousarray(np.argwhere(one_i).flatten(), np.uint32))
            slow_reduction_transfer(
                np.ascontiguousarray(np.argwhere(one_i).flatten(), np.uint32),
                j, i_index, i_count, x, u, v, costs)
        #
        # Perform augmenting row reduction on unassigned i
        #
        ii = np.ascontiguousarray(np.argwhere(free_i).flatten(), np.uint32)
        if len(ii) > 0:
            for iii in range(augmenting_row_reductions):
                ii = slow_augmenting_row_reduction(
                    n, ii, j, i_index, i_count, x, y, u, v, costs)
        slow_augment(n, ii, j, i_index, i_count, x, y, u, v, costs)
    else:
        if np.any(one_i):
            reduction_transfer(
                np.ascontiguousarray(np.argwhere(one_i).flatten(), np.uint32),
                j, i_index, i_count, x, u, v, costs)
        #
        # Perform augmenting row reduction on unassigned i
        #
        ii = np.ascontiguousarray(np.argwhere(free_i).flatten(), np.uint32)
        if len(ii) > 0:
            for iii in range(augmenting_row_reductions):
                ii = augmenting_row_reduction(n, ii, j, i_index, i_count, x, y, u, v, costs)
        augment(n, ii,
                j, i_index, i_count, x, y, u, v, costs)
    if wants_dual_variables:
        return x, y, u, v
    else:
        return x, y


def slow_reduction_transfer(ii, j, idx, count, x, u, v, c):  # pragma: no cover
    '''Perform the reduction transfer step from the Jonker-Volgenant algorithm

    The data is input in a ragged array in terms of "i" structured as a
    vector of values for each i,j combination where:

    ii - the i to be reduced
    j - the j-index of every entry
    idx - the index of the first entry for each i
    count - the number of entries for each i
    x - the assignment of j to i
    u - the dual variable "u" which will be updated. It should be
        initialized to zero for the first reduction transfer.
    v - the dual variable "v" which will be reduced in-place
    c - the cost for each entry.

    The code described in the paper is:

    for each assigned row i do
    begin
       j1:=x[i]; u=min {c[i,j]-v[j] | j=1..n, j != j1};
       v[j1]:=v[j1]-(u-u[i]);
       u[i] = u;
    end;

    The authors note that reduction transfer can be applied in later stages
    of the algorithm but does not seem to provide a substantial benefit
    in speed.
    '''
    for i in ii:
        j1 = x[i]
        jj = j[idx[i]:(idx[i]+count[i])]
        try:
            uu = np.min((c[idx[i]:(idx[i]+count[i])] - v[jj])[jj != j1])
        except ValueError:
            print(c[idx[i]:(idx[i] + count[i])] - v[jj])
            print(idx[i], idx[i] + count[i], jj, j1)
            continue  # raise
        v[j1] = v[j1] - uu + u[i]
        u[i] = uu


def slow_augmenting_row_reduction(n, ii, jj, idx, count, x, y, u, v, c):  # pragma: no cover
    '''Perform the augmenting row reduction step
    from the Jonker-Volgenaut algorithm

    n - the number of i and j in the linear assignment problem
    ii - the unassigned i
    jj - the j-index of every entry in c
    idx - the index of the first entry for each i
    count - the number of entries for each i
    x - the assignment of j to i
    y - the assignment of i to j
    u - the dual variable "u" which will be updated. It should be
        initialized to zero for the first reduction transfer.
    v - the dual variable "v" which will be reduced in-place
    c - the cost for each entry.

    returns the new unassigned i
    '''

    #######################################
    #
    # From Jonker:
    #
    # procedure AUGMENTING ROW REDUCTION;
    # begin
    # LIST: = {all unassigned rows};
    # for all i in LIST do
    #    repeat
    #    ul:=min {c[i,j]-v[j] for j=l ...n};
    #    select j1 with c [i,j1] - v[j1] = u1;
    #    u2:=min {c[i,j]-v[j] for j=l ...n,j< >jl} ;
    #    select j2 with c [i,j2] - v [j2] = u2 and j2 < >j 1 ;
    #    u[i]:=u2;
    #    if u1 <u2 then v[j1]:=v[j1]-(u2-ul)
    #    else if jl is assigned then jl : =j2;
    #    k:=y [jl]; if k>0 then x [k]:=0; x[i]:=jl; y [ j l ] : = i ; i:=k
    #  until ul =u2 (* no reduction transfer *) or k=0 i~* augmentation *)
    #  end
    ii = list(ii)
    k = 0
    limit = len(ii)
    free = []
    while k < limit:
        i = ii[k]
        k += 1
        j = jj[idx[i]:(idx[i] + count[i])]
        uu = c[idx[i]:(idx[i] + count[i])] - v[j]
        order = np.lexsort([uu])
        u1, u2 = uu[order[:2]]
        j1, j2 = j[order[:2]]
        i1 = y[j1]
        if u1 < u2:
            v[j1] = v[j1] - u2 + u1
        elif i1 != n:
            j1 = j2
            i1 = y[j1]
        if i1 != n:
            if u1 < u2:
                k -= 1
                ii[k] = i1
            else:
                free.append(i1)
        x[i] = j1
        y[j1] = i
    return np.array(free, np.uint32)


def slow_augment(n, ii, jj, idx, count, x, y, u, v, c):  # pragma: no cover
    '''Perform the augmentation step to assign unassigned i and j

    n - the # of i and j, also the marker of unassigned x and y
    ii - the unassigned i
    jj - the ragged arrays of j for each i
    idx - the index of the first j for each i
    count - the number of j for each i
    x - the assignments of j for each i
    y - the assignments of i for each j
    u,v - the dual variables
    c - the costs
    '''

    ##################################################
    #
    # Augment procedure: from the Jonker paper.
    #
    # Note:
    #    cred [i,j] = c [i,j] - u [i] - v[j]
    #
    # procedure AUGMENT;
    # begin
    #   for all unassigned i* do
    #   begin
    #     for j:= 1 ... n do
    #       begin d[j] := c[i*,j] - v[j] ; pred[j] := i* end;
    #     READY: = { ) ; SCAN: = { } ; TODO: = { 1 ... n} ;
    #     repeat
    #       if SCAN = { } then
    #       begin
    #         u = min {d[j] for j in TODO} ;
    #         SCAN: = {j | d[j] = u} ;
    #         TODO: = TODO - SCAN;
    #         for j in SCAN do if y[j]==0 then go to augment
    #       end;
    #       select any j* in SCAN;
    #       i := y[j*]; SCAN := SCAN - {j*} ; READY: = READY + {j*} ;
    #       for all j in TODO do if u + cred[i,j] < d[j] then
    #       begin
    #         d[j] := u + cred[i,j]; pred[j] := i;
    #         if d[j] = u then
    #           if y[j] is unassigned then go to augment else
    #           begin SCAN: = SCAN + {j} ; TODO: = TODO - {j} end
    #       end
    #    until false; (* repeat always ends with go to augment *)
    #augment:
    #   (* price updating *)
    #   for k in READY do v[k]: = v[k] + d[k] - u;
    #   (* augmentation *)
    #   repeat
    #     i: = pred[j]; y[ j ] := i ; k:=j; j:=x[i]; x[i]:= k
    #   until i = i*
    #  end
    #end
    inf = np.sum(c) + 1
    d = np.zeros(n)
    cc = np.zeros((n, n))
    cc[:, :] = inf
    for i in range(n):
        cc[i, jj[idx[i]:(idx[i]+count[i])]] = c[idx[i]:(idx[i] + count[i])]
    c = cc
    for i in ii:
        print("Processing i=%d" % i)
        j = jj[idx[i]:(idx[i] + count[i])]
        d = c[i, :] - v
        pred = np.ones(n, int) * i
        on_deck = []
        ready = []
        scan = []
        to_do = list(range(n))
        try:
            while True:
                print("Evaluating i=%d, n_scan = %d" % (i, len(scan)))
                if len(scan) == 0:
                    ready += on_deck
                    on_deck = []
                    umin = np.min([d[jjj] for jjj in to_do])
                    print("umin = %f" % umin)
                    scan = [jjj for jjj in to_do if d[jjj] == umin]
                    to_do = [jjj for jjj in to_do if d[jjj] != umin]
                    for j1 in scan:
                        if y[j1] == n:
                            raise StopIteration()
                j1 = scan[0]
                iii = y[j1]
                print("Consider replacing i=%d, j=%d" % (iii, j1))
                scan = scan[1:]
                on_deck += [j1]
                u1 = c[iii, j1] - v[j1] - umin
                for j1 in list(to_do):
                    h = c[iii, j1] - v[j1] - u1
                    print("Consider j=%d as replacement, c[%d,%d]=%f,v[%d]=%f,h=%f, d[j]= %f" %
                          (j1, iii, j1, c[iii, j1], j1, v[j1], h, d[j1]))
                    if h < d[j1]:
                        print("Add to chain")
                        pred[j1] = iii
                        if h == umin:
                            if y[j1] == n:
                                raise StopIteration()
                            print("Add to scan")
                            scan += [j1]
                            to_do.remove(j1)
                        d[j1] = h

        except StopIteration:
            # Augment
            print("Augmenting %d" % j1)
            for k in ready:
                temp = v[k]
                v[k] = v[k] + d[k] - umin
                print("v[%d] %f -> %f" % (k, temp, v[k]))
            while True:
                iii = pred[j1]
                print("y[%d] %d -> %d" % (j1, y[j1], iii))
                y[j1] = iii
                j1, x[iii] = x[iii], j1
                if iii == i:
                    break
    #
    # Re-establish slackness since we didn't pay attention to u
    #
    for i in range(n):
        j = x[i]
        u[i] = c[i, j] - v[j]

# if __name__=="__main__":
#     i = np.load("c:/temp/bad/img-1557/i.npy")
#     j = np.load("c:/temp/bad/img-1557/j.npy")
#     costs = np.load("c:/temp/bad/img-1557/c.npy")
#     lapjv(i, j, costs)
