import model
import loss
import optim

# Inputs:
#     train = (X_train, Y_train):
#         X_train: A n_train * m array, the ith instance is stored in X_train[i, : ]
#         Y_train: A n_train * l array, if the ith instance belongs to the jth class, then Y_train[i, j] = 1, otherwise Y_train[i, j] = 0
#     test = (X_test, Y_test):
#         X_test: A n_test * m array, the ith instance is stored in X_test[i, : ]
#         Y_test: A n_test * l array, if the ith instance belongs to the jth class, then Y_test[i, j] = 1, otherwise Y_test[i, j] = 0    
#     alpha: 
#     lambda1: 
#     lambda2: 
#     lambda3: 
#
# Outputs:
#     P: A n_test * l array, if the ith instance is predicted to belong to the jth class, then P[i, j] = 1, otherwise P[i, j] = 0
#     O: A n_test * l array, O[i, : ] is the output of the model for the ith instance
def run(train, test, alpha, lambda1, lambda2, lambda3):
    X, Y = train

    m, l = X.shape[1], Y.shape[1]

    linear = model.Linear(X, Y, round(alpha * min(m, l)))
    wrap = loss.Wrap(linear, lambda1, lambda2, lambda3)

    alternate = optim.Alternate(linear, wrap, 0.1)
    alternate.optimize()

    X, Y = test

    return linear.predict(X), linear.output(X)
