# Path: pyglm/inference/coord_descent.py
# def fit_glm(xn, n, 
#             (glm_syms, glm_nll, g_glm_nll)):
#     """ Fit the GLM parameters in state dict x
#     """
#     # Get the differentiable variables for the n-th GLM
#     dnvars = get_vars(glm_syms, xn['glm'])
#     x_glm_0, shapes = packdict(dnvars)
#     
#     # Create lambda functions to compute the nll and its gradient and Hessian
#     def nll(x_glm_vec):
#         y = glm_nll(x_glm_vec, xn)
#         if np.isnan(y):
#             y = 1e16
# 
#         return y
# 
#     def grad_nll(x_glm_vec):
#         g = g_glm_nll(x_glm_vec, xn)
#         if np.any(np.isnan(g)):
#             g = np.zeros_like(g)
# 
#         return g
#     
#     # Callback to print progress. In order to count iters, we need to
#     # pass the current iteration via a list
#     ncg_iter_ls = [0]
#     def progress_report(x_curr, ncg_iter_ls):
#         ll = -1.0*nll(x_curr)
#         print "Newton iter %d.\tNeuron %d. LL: %.1f" % (ncg_iter_ls[0],n,ll)
#         ncg_iter_ls[0] += 1
#     cbk = lambda x_curr: progress_report(x_curr, ncg_iter_ls)
# 
#     # Call the appropriate scipy optimization function
#     res = opt.minimize(nll, x_glm_0,
#                        method="bfgs",
#                        jac=grad_nll,
#                        options={'disp': True,
#                                 'maxiter' : 225},
#                        callback=cbk)
#     xn_opt = res.x
# 
#     # Unpack the optimized parameters back into the state dict
#     x_glm_n = unpackdict(xn_opt, shapes)
#     set_vars(glm_syms, xn['glm'], x_glm_n)
# 
# Path: pyglm/utils/progress_report.py
# def wait_watching_stdout(ar, interval=1, truncate=100):
#     """ Print the outputs of each worker 
#     """
#     stdoffs = [0]*len(ar.stdout)
#     while not ar.ready():
#         stdouts = ar.stdout
#         if not any(stdouts):
#             continue
# 
#         for eid, stdout in enumerate(stdouts):
#             if stdout:
#                 newoff = len(stdout)
#                 if newoff > stdoffs[eid]:
#                     print "[ stdout %2i ]\n%s" % (eid, stdout[stdoffs[eid]:])
#                     stdoffs[eid] = newoff
# 
#         print '-' * 30
#         print "%.3fs elapsed. Progress: %d/%d" % (ar.elapsed, ar.progress, len(ar))
#         print ""
# 
#         sys.stdout.flush()
#         time.sleep(interval)


import numpy as np
from IPython.parallel.util import interactive
from pyglm.inference.coord_descent import fit_glm
from pyglm.utils.progress_report import wait_watching_stdout
                  block=True)

    # Parallel function to fit GLMs
    @interactive
    def _parallel_fit_glm(n, x, use_hessian=False, use_rop=False):
        nvars = popn.extract_vars(x, n)
        fit_glm(nvars, n, glm_inf_prms, use_hessian, use_rop)
        return nvars['glm']

    # Alternate fitting the network and fitting the GLMs
    x = x0
    lp_prev = parallel_compute_log_p(dview, master, x, N)
    converged = False
    iter = 0
    while not converged and iter < maxiter:
        iter += 1
        print "Coordinate descent iteration %d." % iter
        
        # TODO Fit the network on the first engine
        # fit_network(x, net_inf_prms, use_hessian, use_rop)
        
        # Fit the GLMs in parallel
        x_glms = dview.map_async(_parallel_fit_glm,
                                 range(N),
                                 [x]*N)

        # Print progress report ever interval seconds
        interval = 15.0
        # Timeout after specified number of seconds (-1 = Inf?)
        #timeout = -1

# Path: pyglm/inference/coord_descent.py
# def fit_glm(xn, n, 
#             (glm_syms, glm_nll, g_glm_nll)):
#     """ Fit the GLM parameters in state dict x
#     """
#     # Get the differentiable variables for the n-th GLM
#     dnvars = get_vars(glm_syms, xn['glm'])
#     x_glm_0, shapes = packdict(dnvars)
#     
#     # Create lambda functions to compute the nll and its gradient and Hessian
#     def nll(x_glm_vec):
#         y = glm_nll(x_glm_vec, xn)
#         if np.isnan(y):
#             y = 1e16
# 
#         return y
# 
#     def grad_nll(x_glm_vec):
#         g = g_glm_nll(x_glm_vec, xn)
#         if np.any(np.isnan(g)):
#             g = np.zeros_like(g)
# 
#         return g
#     
#     # Callback to print progress. In order to count iters, we need to
#     # pass the current iteration via a list
#     ncg_iter_ls = [0]
#     def progress_report(x_curr, ncg_iter_ls):
#         ll = -1.0*nll(x_curr)
#         print "Newton iter %d.\tNeuron %d. LL: %.1f" % (ncg_iter_ls[0],n,ll)
#         ncg_iter_ls[0] += 1
#     cbk = lambda x_curr: progress_report(x_curr, ncg_iter_ls)
# 
#     # Call the appropriate scipy optimization function
#     res = opt.minimize(nll, x_glm_0,
#                        method="bfgs",
#                        jac=grad_nll,
#                        options={'disp': True,
#                                 'maxiter' : 225},
#                        callback=cbk)
#     xn_opt = res.x
# 
#     # Unpack the optimized parameters back into the state dict
#     x_glm_n = unpackdict(xn_opt, shapes)
#     set_vars(glm_syms, xn['glm'], x_glm_n)
# 
# Path: pyglm/utils/progress_report.py
# def wait_watching_stdout(ar, interval=1, truncate=100):
#     """ Print the outputs of each worker 
#     """
#     stdoffs = [0]*len(ar.stdout)
#     while not ar.ready():
#         stdouts = ar.stdout
#         if not any(stdouts):
#             continue
# 
#         for eid, stdout in enumerate(stdouts):
#             if stdout:
#                 newoff = len(stdout)
#                 if newoff > stdoffs[eid]:
#                     print "[ stdout %2i ]\n%s" % (eid, stdout[stdoffs[eid]:])
#                     stdoffs[eid] = newoff
# 
#         print '-' * 30
#         print "%.3fs elapsed. Progress: %d/%d" % (ar.elapsed, ar.progress, len(ar))
#         print ""
# 
#         sys.stdout.flush()
#         time.sleep(interval)


# Path: pyglm/inference/parallel_coord_descent.py
import numpy as np
from IPython.parallel.util import interactive
from pyglm.inference.coord_descent import fit_glm
from pyglm.utils.progress_report import wait_watching_stdout

                  block=True)

    # Parallel function to fit GLMs
    @interactive
    def _parallel_fit_glm(n, x, use_hessian=False, use_rop=False):
        nvars = popn.extract_vars(x, n)
        fit_glm(nvars, n, glm_inf_prms, use_hessian, use_rop)
        return nvars['glm']

    # Alternate fitting the network and fitting the GLMs
    x = x0
    lp_prev = parallel_compute_log_p(dview, master, x, N)
    converged = False
    iter = 0
    while not converged and iter < maxiter:
        iter += 1
        print "Coordinate descent iteration %d." % iter
        
        # TODO Fit the network on the first engine
        # fit_network(x, net_inf_prms, use_hessian, use_rop)
        
        # Fit the GLMs in parallel
        x_glms = dview.map_async(_parallel_fit_glm,
                                 range(N),
                                 [x]*N)

        # Print progress report ever interval seconds
        interval = 15.0
        # Timeout after specified number of seconds (-1 = Inf?)
        #timeout = -1

        wait_watching_stdout(x_glms, interval=interval)