# Path: graph_utils.py
# def decompose_graph(X):
#     contains = []
#     chains = []
#     edge_index = []
# 
#     for x in X:
#         _contains, _chains, _edge_index = monotonic_chains(x[0].shape[0], x[1])
#         contains.append(_contains)
#         chains.append(_chains)
#         edge_index.append(_edge_index)
# 
#     return contains, chains, edge_index
# 
# def decompose_grid_graph(X, get_sign=False):
#     contains_node = []
#     chains = []
#     edge_index = []
#     sign = []
# 
#     n_states = 10
#     n_nodes = 400
#     width = 20
#     height = 20
# 
#     for k in xrange(len(X)):
#         x = X[k]
#         _edge_index = {}
#         for i, edge in enumerate(x[1]):
#             _edge_index[(edge[0], edge[1])] = i
# 
#         _sign = []
#         _chains = []
#         _contains = [[] for i in xrange(n_nodes)]
# 
#         #vertical
#         for i in xrange(0, n_nodes, width):
#             _sign.append(1)
#             _chains.append(np.arange(i, i + width))
#             assert _chains[-1].shape[0] == width
#             tree_number = len(_chains) - 1
#             for node in _chains[-1]:
#                 _contains[node].append(tree_number)
# 
#         #horizontal
#         for i in xrange(0, width):
#             _sign.append(-1)
#             _chains.append(np.arange(i, n_nodes, width))
#             assert _chains[-1].shape[0] == height
#             tree_number = len(_chains) - 1
#             for node in _chains[-1]:
#                 _contains[node].append(tree_number)
# 
#         sign.append(_sign)
#         contains_node.append(_contains)
#         chains.append(_chains)
#         edge_index.append(_edge_index)
# 
#     if get_sign:
#         return contains_node, chains, edge_index, sign
#     else:
#         return contains_node, chains, edge_index


import numpy as np
from scipy.optimize import fmin_l_bfgs_b
from graph_utils import decompose_graph, decompose_grid_graph
from trw_utils import *


def f(x, node_weights, multiplier, chains, edge_weights, edge_index, sign, y_hat, contains_node, history):
    dual = 0
    n_nodes, n_states = node_weights.shape
    unaries = node_weights * multiplier
    lambdas = x.reshape((400, 10))

    for i, chain in enumerate(chains):
        y_hat[i], e = optimize_chain(chain,
                                     sign[i] * lambdas[chain,:] + unaries[chain,:],
                                     edge_weights,
                                     edge_index)

        dual += e

    dlambda = np.zeros((n_nodes, n_states), dtype=np.float64)
    for p in xrange(n_nodes):
        for i in contains_node[p]:
            pos = np.where(chains[i] == p)[0][0]
            dlambda[p][y_hat[i][pos]] += sign[i]

    history.append(dual)

    return dual, dlambda.reshape(4000)


def trw_lbfgs(node_weights, edges, edge_weights,
              max_iter=100, verbose=1, tol=1e-3):


# Path: graph_utils.py
# def decompose_graph(X):
#     contains = []
#     chains = []
#     edge_index = []
# 
#     for x in X:
#         _contains, _chains, _edge_index = monotonic_chains(x[0].shape[0], x[1])
#         contains.append(_contains)
#         chains.append(_chains)
#         edge_index.append(_edge_index)
# 
#     return contains, chains, edge_index
# 
# def decompose_grid_graph(X, get_sign=False):
#     contains_node = []
#     chains = []
#     edge_index = []
#     sign = []
# 
#     n_states = 10
#     n_nodes = 400
#     width = 20
#     height = 20
# 
#     for k in xrange(len(X)):
#         x = X[k]
#         _edge_index = {}
#         for i, edge in enumerate(x[1]):
#             _edge_index[(edge[0], edge[1])] = i
# 
#         _sign = []
#         _chains = []
#         _contains = [[] for i in xrange(n_nodes)]
# 
#         #vertical
#         for i in xrange(0, n_nodes, width):
#             _sign.append(1)
#             _chains.append(np.arange(i, i + width))
#             assert _chains[-1].shape[0] == width
#             tree_number = len(_chains) - 1
#             for node in _chains[-1]:
#                 _contains[node].append(tree_number)
# 
#         #horizontal
#         for i in xrange(0, width):
#             _sign.append(-1)
#             _chains.append(np.arange(i, n_nodes, width))
#             assert _chains[-1].shape[0] == height
#             tree_number = len(_chains) - 1
#             for node in _chains[-1]:
#                 _contains[node].append(tree_number)
# 
#         sign.append(_sign)
#         contains_node.append(_contains)
#         chains.append(_chains)
#         edge_index.append(_edge_index)
# 
#     if get_sign:
#         return contains_node, chains, edge_index, sign
#     else:
#         return contains_node, chains, edge_index


# Path: trw.py
import numpy as np
from scipy.optimize import fmin_l_bfgs_b
from graph_utils import decompose_graph, decompose_grid_graph
from trw_utils import *



def f(x, node_weights, multiplier, chains, edge_weights, edge_index, sign, y_hat, contains_node, history):
    dual = 0
    n_nodes, n_states = node_weights.shape
    unaries = node_weights * multiplier
    lambdas = x.reshape((400, 10))

    for i, chain in enumerate(chains):
        y_hat[i], e = optimize_chain(chain,
                                     sign[i] * lambdas[chain,:] + unaries[chain,:],
                                     edge_weights,
                                     edge_index)

        dual += e

    dlambda = np.zeros((n_nodes, n_states), dtype=np.float64)
    for p in xrange(n_nodes):
        for i in contains_node[p]:
            pos = np.where(chains[i] == p)[0][0]
            dlambda[p][y_hat[i][pos]] += sign[i]

    history.append(dual)

    return dual, dlambda.reshape(4000)


def trw_lbfgs(node_weights, edges, edge_weights,
              max_iter=100, verbose=1, tol=1e-3):


    result = decompose_grid_graph([(node_weights, edges, edge_weights)], get_sign=True)