#!/usr/bin/env python
# -*- coding: utf-8 -*-
# =======================================
# File Name: ADMM_primal.py
# Purpose  : implementation for ADMM method
#            for solving primal problem
# =======================================

from utils import get_params
import numpy as np
import sys


def ADMM_primal(mu, nu, c, iters=10000, rho=1024, alpha=1.618):
    """ADMM_primal
    """
    # initialize
    m, n = c.shape
    pi = np.zeros((m, n))
    pi_dag = np.zeros((m, n))
    w = np.zeros((m, n))
    u = np.zeros(m)
    v = np.zeros(n)

    rho_tilde = rho * 32
    while rho_tilde >= rho:
        for _ in range(iters):
            r = ((-w + u.reshape((m, 1)) + v.reshape((1, n)) - c) / rho + 
                mu.reshape((m, 1)) + nu.reshape((1, n)) + pi_dag)
        
            pi = (r - ((r.sum(axis=1) - r.sum() / (m + n + 1)) / (n + 1)).reshape((m, 1))
                - ((r.sum(axis=0) - r.sum() / (m + n + 1)) / (m + 1)).reshape((1, n)))

            pi_dag = np.maximum(pi + w / rho, 0.0)

            u = u + alpha * rho * (mu - pi.sum(axis=1))
            v = v + alpha * rho * (nu - pi.sum(axis=0))
            w = w + alpha * rho * (pi - pi_dag)

            rho_tilde = rho_tilde / 2

        print('error_mu = %.5e' % np.linalg.norm(pi_dag.sum(axis = 1) - mu, 1))
        print('error_nu = %.5e' % np.linalg.norm(pi_dag.sum(axis = 0) - nu, 1))
        print('fvall    = %.5e' % (c * pi_dag).sum())


if __name__ == '__main__':
    try:
        print("Test...")
        _mu, _nu, _c = get_params(64, 'random')
        ADMM_primal(_mu, _nu, _c)
    except KeyboardInterrupt:
        print ("  Ctrl+C pressed...")
        sys.exit(1)


[((19, 9, 19, 25), 'numpy.zeros', 'np.zeros', ({(19, 18, 19, 24): '(m, n)'}, {}), '((m, n))', True, 'import numpy as np\n'), ((20, 13, 20, 29), 'numpy.zeros', 'np.zeros', ({(20, 22, 20, 28): '(m, n)'}, {}), '((m, n))', True, 'import numpy as np\n'), ((21, 8, 21, 24), 'numpy.zeros', 'np.zeros', ({(21, 17, 21, 23): '(m, n)'}, {}), '((m, n))', True, 'import numpy as np\n'), ((22, 8, 22, 19), 'numpy.zeros', 'np.zeros', ({(22, 17, 22, 18): 'm'}, {}), '(m)', True, 'import numpy as np\n'), ((23, 8, 23, 19), 'numpy.zeros', 'np.zeros', ({(23, 17, 23, 18): 'n'}, {}), '(n)', True, 'import numpy as np\n'), ((50, 23, 50, 47), 'utils.get_params', 'get_params', ({(50, 34, 50, 36): '64', (50, 38, 50, 46): '"""random"""'}, {}), "(64, 'random')", False, 'from utils import get_params\n'), ((34, 21, 34, 50), 'numpy.maximum', 'np.maximum', ({(34, 32, 34, 44): 'pi + w / rho', (34, 46, 34, 49): '0.0'}, {}), '(pi + w / rho, 0.0)', True, 'import numpy as np\n'), ((54, 8, 54, 19), 'sys.exit', 'sys.exit', ({(54, 17, 54, 18): '(1)'}, {}), '(1)', False, 'import sys\n')]