import theano
import cascade_creation
from theano import tensor, function
import numpy as np
import timeout
import cvxopt


@timeout.timeout(20)
def sparse_recovery(lbda, n_cascades):
    """
    Solves:
        min lbda * |theta|_1 - 1/n_cascades*(
                sum w log(e^[M*theta]) + (1-w) log(1-e^[M*theta])
            )
        s.t theta_j <= 0

    The objective must be re-normalized:
    -> log(e^x - 1) = (1-k)x + log(e^[kx] - 1)
    with k chosen as : n_nodes*n_cascades
    """
    if type(lbda) == int or type(lbda) == float:
        lbda = np.array(lbda)

    lbda = theano.shared(lbda.astype(theano.config.floatX))

    theta = tensor.row().T
    theta_ = theta.flatten()

    M = tensor.matrix(name='M', dtype=theano.config.floatX)
    w = tensor.vector(name='w', dtype=theano.config.floatX)

    y =  lbda * theta_.norm(1) - 1./n_cascades*(
            (w).dot(tensor.log(1-tensor.exp(M.dot(theta_))))
            + (1-w).dot(tensor.dot(M, theta_))
        )  

    z = tensor.row().T
    z_ = z.flatten()

    y_diff = tensor.grad(y, theta_)
    y_hess = z[0] * theano.gradient.hessian(y, theta_)
    f_x = theano.function([theta, M, w], [y, y_diff],
                        allow_input_downcast=True)
    f_xz = theano.function([theta, z, M, w], [y, y_diff, y_hess],
                        allow_input_downcast=True)

    return f_x, f_xz


def type_lasso(lbda):
    """
    Solves:
        min - sum_j theta_j + lbda*|e^{M*theta} - (1 - w)|_2
        s.t theta_j <= 0
    """
    if type(lbda) == int or type(lbda) == float:
        lbda = np.array(lbda)

    lbda = theano.shared(lbda.astype(theano.config.floatX))

    theta = tensor.row().T
    theta_ = theta.flatten()

    M = tensor.matrix(name='M', dtype=theano.config.floatX)
    w = tensor.vector(name='w', dtype=theano.config.floatX)

    y = lbda * (theta_).norm(1) + (
        tensor.exp(M.dot(theta_)) - (1 - w)).norm(2)

    z = tensor.row().T
    z_ = z.flatten()

    y_diff = tensor.grad(y, theta_)
    y_hess = z[0] * theano.gradient.hessian(y, theta_)
    f_x = theano.function([theta, M, w], [y, y_diff],
                        allow_input_downcast=True)
    f_xz = theano.function([theta, z, M, w], [y, y_diff, y_hess],
                        allow_input_downcast=True)

    return f_x, f_xz

@timeout.timeout(70)
def diff_and_opt(M_val, w_val, f_x, f_xz):

    if M_val.dtype == bool:
        M_val = M_val.astype(theano.config.floatX)

    m, n = M_val.shape

    def F(x=None, z=None):
        if x is None:
            return 0, cvxopt.matrix(-.1, (n,1))
        elif z is None:
            y, y_diff = f_x(x, M_val, w_val)
            return cvxopt.matrix(float(y), (1, 1)),\
                   cvxopt.matrix(y_diff.astype("float64")).T
        else:
            y, y_diff, y_hess = f_xz(x, z, M_val, w_val)
            return cvxopt.matrix(float(y), (1, 1)), \
                    cvxopt.matrix(y_diff.astype("float64")).T, \
                    cvxopt.matrix(y_hess.astype("float64"))

    G = cvxopt.spdiag([1 for i in range(n)])
    h = cvxopt.matrix(0.0, (n,1))

    #Relaxing precision constraints
    # cvxopt.solvers.options['feastol'] = 2e-5
    # cvxopt.solvers.options['abstol'] = 2e-5
    # cvxopt.solvers.options['maxiters'] = 50
    cvxopt.solvers.options['show_progress'] = False
    try:
        theta = cvxopt.solvers.cp(F, G, h)['x']
    except ArithmeticError:
        print("ArithmeticError thrown, change initial point"
              " given to the solver")
    except ValueError:
        print("Domain Error, skipping to next node")
        theta = np.zeros(len(w_val))

    if cvxopt.solvers.options['show_progress']:
        print(1 - np.exp(theta))

    return 1 - np.exp(theta), theta


def test():
    """
    unit test
    """
    lbda = .0001
    G = cascade_creation.InfluenceGraph(max_proba=.9)
    G.erdos_init(n=20, p = .3)
    A = cascade_creation.generate_cascades(G, .1, 500)
    M_val, w_val = cascade_creation.icc_matrixvector_for_node(A, 2)
    print(len(M_val))
    
    #Type lasso
    if 0:
        f_x, f_xz = type_lasso(lbda)
        p_vec, _ = diff_and_opt(M_val, w_val, f_x, f_xz)
        print(G.mat[2])

    #Sparse recovery
    if 1:
        f_x, f_xz = sparse_recovery(lbda, 500)
        p_vec, _ = diff_and_opt(M_val, w_val, f_x, f_xz)
        print(G.mat[2])

if __name__=="__main__":
    test()