"""
Problem 2: Two-layer neural network trained by gradient descent.
Usage: python main.py <train.dat> <test.dat>
Output: accuracy on stderr; predictions written to {test_filename}_sol.dat in same form as test.dat.
"""

import os
import numpy as np


def load_data(path):
    data = np.loadtxt(path)
    X = data[:, :2]
    y = data[:, 2].astype(np.int64)
    return X, y


# Do not change this function signature (different values will be used for autograding)
# but use different hyperparameter values for your experiments!
def main(
    train_path: str,
    test_path: str,
    m: int = None,
    eta: float = None,
    tol: float = 1e-5,
    maxiter: int = None,
):
    X_train, y_train = load_data(train_path)
    X_test, y_test = load_data(test_path)

    # Do some machine learning here
    preds = (
        np.arange(len(y_test)) % 2
    )  # hopefully your prediction will be better than this

    out_path = os.path.join(
        os.path.dirname(test_path),
        os.path.splitext(os.path.basename(test_path))[0] + "-pred.dat",
    )
    with open(out_path, "w") as f:
        for i in range(len(X_test)):
            f.write(f"{X_test[i, 0]:.6f} {X_test[i, 1]:.6f} {preds[i]}\n")


if __name__ == "__main__":
    main("hw5-nn-train-100.dat", "hw5-nn-test.dat")