from IPython.display import clear_output
import time

def custom_gridsearch(trainer, grid, train_data, valid_data, minimize=True):
    valid_scores = []

    time_start = time.time()
    evaluator = RmseEvaluator(labelCol="rating", predictionCol="prediction")

    n = len(grid)
    for i, params in enumerate(grid):
        for k,v in params.items():
            k.parent = v
        model = trainer.fit(train_data)
        predictions = model.transform(valid_data)
        valid_score = evaluator.evaluate(predictions)
        valid_scores.append(valid_score)
        # tqdm not working here
        clear_output()
        print("valid_score: {}, time_left: {}".format(valid_score, (time.time() - time_start) * (n-i-1)/(i+1)))
    print("time passed: {}".format(time.time()-time_start))

    if minimize:
        best_idx = np.argmin(valid_scores)
    else:
        best_idx = np.argmax(valid_scores)
    # return best_params, best_score,
    return grid[best_idx], valid_scores[best_idx], valid_scores

best_params, best_score, all_scores = custom_gridsearch(als, grid, dev, val)