h2o.init()

from h2o.utils.shared_utils import _locate # private function. used to find files within h2o git project directory.

# Import walking gait data
gait = h2o.upload_file(path=_locate("smalldata/glrm_test/subject01_walk1.csv"))
gait.describe()

# Basic GLRM using quadratic loss and no regularization (PCA)
model = H2OGeneralizedLowRankEstimator(k=10, loss="Quadratic", regularization_x="None", regularization_y="None", max_iterations=1000)
model.train(x=list(range(1,gait.ncol)), training_frame=gait)
model.show()

# Plot objective function value each iteration
model_score = model.score_history()
plt.xlabel("Iteration")
plt.ylabel("Objective")
plt.title("Objective Function Value per Iteration")
print(model_score)
plt.plot(model_score["iteration"], model_score["objective"])
plt.show()

# Archetype to feature mapping (Y)
gait_y = model._model_json["output"]["archetypes"]
print(gait_y)

gait_y_np = np.array(model.archetypes())
x_coords = range(0, gait_y_np.shape[1], 3)
y_coords = range(1, gait_y_np.shape[1], 3)

# Plot archetypes on x- vs. y-coordinate features
for k in range(gait_y_np.shape[0]):
    x_pts = gait_y_np[k, x_coords]
    y_pts = gait_y_np[k, y_coords]
    plt.plot(x_pts, y_pts, 'bo')

    # Add feature labels to each point
    feat_names = [nam[:-1] for nam in gait_y.col_header[1::3]]
    for i in range(len(feat_names)):
        plt.annotate(feat_names[i], xy = [x_pts[i], y_pts[i]])
    plt.title("Feature Weights of Archetype " + str(k+1))
    plt.xlabel("X-Coordinate Weight")
    plt.ylabel("Y-Coordinate Weight")
    plt.show()

# Projection into archetype space (X)
x_key = model._model_json["output"]["representation_name"]
gait_x = h2o.get_frame(x_key)
gait_x.show()

time_np = np.array(h2o.as_list(gait["Time"]))
gait_x_np = np.array(h2o.as_list(gait_x))

# Plot archetypes over time
lines = []
for i in range(gait_x_np.shape[1]):
    lines += plt.plot(time_np, gait_x_np[:,i], '-')
plt.title("Archetypes over Time")
plt.xlabel("Time")
plt.ylabel("Archetypal Projection")
plt.legend(lines, gait_x.col_names)
plt.show()

# Reconstruct data from X and Y
pred = model.predict(gait)
pred.head()

# Import walking gait data with missing values
gait_miss = h2o.upload_file(_locate("smalldata/glrm_test/subject01_walk1_miss15.csv"))
gait_miss.describe()