import numpy as np
import time
from scipy.optimize import minimize, nnls
from cvxpy import *

np.random.seed(1)

""" Test-data generator """
def create_test_data(samples, features, noise, loc=10, scale=2.5):
    m = np.random.normal(loc=loc, scale=scale, size=(samples, features))
    x = np.abs(np.random.normal(size=m.shape[0]))
    y = np.dot(x, m)
    y += np.random.normal(loc=0, scale=noise)
    return np.clip(m, 0, np.inf), y

""" SLSQP-based approach """
def solve_slsqp(m, y):
    def loss(x):
        return np.sum(np.square((np.dot(x, m) - y)))

    cons = ({'type': 'eq',
             'fun' : lambda x: np.sum(x) - 1.0})

    x0 = np.zeros(m.shape[0])
    start = time.time()
    res = minimize(loss, x0, method='SLSQP', constraints=cons,
                   bounds=[(0, np.inf) for i in range(m.shape[0])], options={'disp': False})
    end = time.time()
    return end-start, res.fun

""" General-purpose SOCP """
def solve_socp(x, y):
    X = Variable(x.shape[0])
    constraints = [X >= 0, sum_entries(X) == 1.0]
    product = x.T * diag(X)
    diff = sum_entries(product, axis=1) - y
    problem = Problem(Minimize(norm(diff)), constraints)
    start = time.time()
    problem.solve()  # only pure solving time is measured!
    end = time.time()
    return end-start, problem.value**2


""" Customized NNLS-based approach """
def solve_nnls(x, y):
    start = time.time()
    B = x.T
    A = B[:, :-1] - B[:, -1:]
    bb, norm = nnls(A, y - B[:, -1])
    bi = np.append(bb, 1 - sum(bb))

    end = time.time()
    return end-start, norm**2

""" Benchmark """

N_ITERS = 5
slsqp_results = []
socp_results = []
nnls_results = []
for i in range(N_ITERS):
    print('it: ', i)
    x, y = create_test_data(20, 100000, 5)

    slsqp_results.append(solve_slsqp(x,y))
    socp_results.append(solve_socp(x,y))
    nnls_results.append(solve_nnls(x,y))

for i in range(N_ITERS):
    print(slsqp_results[i])
    print(socp_results[i])
    print(nnls_results[i])


print('avg(slsqp): ', sum(map(lambda x: x[0], slsqp_results)))
print('avg(socp): ', sum(map(lambda x: x[0], socp_results)))
print('avg(nnls): ', sum(map(lambda x: x[0], nnls_results)))

# it:  0
# it:  1
# it:  2
# it:  3
# it:  4
# (6.922792911529541, 1209046597.8513827)
# (15.238645076751709, 2739096722.785817)
# (0.0748591423034668, 2737550283.250587)
# (5.709033966064453, 1222823702.2829199)
# (11.623920917510986, 2250555718.285656)
# (0.06250286102294922, 2249893123.793492)
# (5.927958011627197, 1176982468.5536709)
# (13.555027961730957, 2338851373.7386045)
# (0.059979915618896484, 2338344485.827066)
# (5.572870969772339, 1086835079.1386988)
# (12.083420038223267, 1963442006.6095839)
# (0.041224002838134766, 1963124553.0984313)
# (5.766523838043213, 1204447559.1084349)
# (12.301261901855469, 2550996627.674802)
# (0.06191396713256836, 2550971465.715446)
# avg(slsqp):  29.899179697036743
# avg(socp):  64.80227589607239
# avg(nnls):  0.3004798889160156